Typed Tensors

Typed tensors provide an alternative API for which tensor characteristics are encoded in the type of the tensor. The Hasktorch library is layered such that [ TODO ]

Using typed tensors increases the expressiveness of program invariants that can be automatically checked by GHC at the cost of needing to be more explicit in writing code and also requiring working with Haskell's type-level machinery. Type-level Haskell programming has arisen through progressive compiler iteration leading to mechanisms that have a higher degree of complexity compared to value-level Haskell code or other languages such as Idris which were designed with type-level computations from inception. In spite of these compromises, Haskell offers enough capability to express powerful type-level representations of models that is matched by few other languages used for production applications.

Things we can do in typed Hasktorch:

specify, check, and infer tensor shapes at compile time
specify, check, and infer tensor data types at compile time
specify, check, and infer tensor compute devices at compile time

We can encode all Hasktorch tensor shapes on the type level using type-level lists and natural numbers:

type EmptyShape = '[]
type OneDimensionalShape (a :: Nat) = '[a]
type TwoDimensionalShape (a :: Nat) (b :: Nat) = '[a, b]
...

Tensor data types and compute device types are lifted to the type level using the DataKinds language extension:

type BooleanCPUTensor (shape :: [Nat]) = Tensor '(CPU,  0) 'Bool  shape
type IntCUDATensor    (shape :: [Nat]) = Tensor '(CUDA, 1) 'Int64 shape

Devices are represented as tuples consisting of a DeviceType (here CPU for the CPU and CUDA for a CUDA device, respectively) and a device id (here the Nats 0 and 1, respectively).

It is a common misconception that specifying tensor properties at compile time is only possible if all tensor properties are constants and are statically known. If this were the case, then we could only write functions over fully specified tensors, say,

boring :: BooleanCPUTensor '[] -> BooleanCPUTensor '[]
boring = id

Fortunately, Haskell has the ability to reason about type variables. This feature is called parametric polymorphism. Consider this simple example of a function:

tensorNoOp
  :: forall (shape :: [Nat]) (dtype :: DType) (device :: (DeviceType, Nat))
   . Tensor device dtype shape
  -> Tensor device dtype shape
tensorNoOp = id

Here, shape, dtype, and device are type variables that have been constrained to be of kind shape, data type, and device, respectively. The universal quantifier forall implies that this function is well-defined for all inhabitants of the types that are compatible with the type variables and their constraints.

The tensorNoOp function may seem trivial, and that is because its type is very strongly constrained: Given any typed tensor, it must return a tensor of the same shape and with the same data type and on the same device. Besides by means of the identity, id, there are not many ways in which this function can be implemented.

There is a connection between a type signature of a function and mathematical proofs. We can say that the fact that a function exists is witnessed by its implementation. The implementation, tensorNoOp = id, is the proof of the theorem stated by tensorNoOp's type signature. And here we have a proof that we can run this function on any device and for any data type and tensor shape.

This is the essence of typed Hasktorch. As soon as the compiler gives its OK, we hold a proof that our program will run without shape or CUDA errors.

Autograd with typed tensors

Chapter 4 introduced automatic differentiation through the untyped API. The typed API exposes the same machinery — the same ATen tape runs underneath — but the types document and enforce the calling convention that the untyped API leaves implicit.

Trainable values are Parameter device dtype shape, the typed counterpart of the untyped IndependentTensor:

w  <- makeIndependent =<< randn @'[2, 1] @'D.Float @'( 'D.CPU, 0)
let w' = toDependent w   -- Tensor '( 'D.CPU, 0) 'D.Float '[2, 1], tracked

Gradients are computed with grad:

grad :: Tensor device dtype '[] -> a -> b

Two things are worth reading off this signature. First, the loss argument has shape '[]: the type demands a scalar. Calling autograd on a non-scalar — a runtime error (or a silent implicit-sum surprise) in dynamic frameworks — does not compile here. Second, grad loss (flattenParameters model) returns an HList of gradient tensors whose shapes are, by type, exactly the shapes of the parameters they belong to. A shape mismatch between a parameter and its gradient update cannot be expressed.

Models are ordinary records of typed layers; deriving Generic and Parameterized gives parameter traversal for free, and HasForward states the shape contract of the forward pass:

data MLP ... = MLP
  { layer0 :: Linear inputFeatures hiddenFeatures dtype device,
    layer1 :: Linear hiddenFeatures outputFeatures dtype device
  }
  deriving (Show, Generic, Parameterized)

instance ... => HasForward (MLP ...)
    (Tensor device dtype '[batchSize, inputFeatures])
    (Tensor device dtype '[batchSize, outputFeatures])
  where
  forward MLP {..} = forward layer1 . tanh . forward layer0

Random initialization comes from Randomizable (sample MLPSpec), and a training step is one call:

(model', optim') <- runStep model optim loss learningRate

where optim is mkGD, mkGDM, or mkAdam, and loss = mseLoss @ReduceMean actual expected is again forced by its type to be a scalar on the right device and dtype. The complete program is in examples/static-xor-mlp.

The end-to-end guarantee is the one this chapter has been building towards: if the program compiles, the forward pass, the loss, every gradient, and every optimizer update agree about shapes, dtype, and device — the whole training loop, not just individual operations.