Operation-Polymorphic Tensor Programs

Program is model code abstract over a tensor-operation interpreter. Supplying the eager interpreter executes operations immediately and records a dynamic tape; supplying the typed-graph recorder, as the autograd transforms do, records reusable shape-indexed SSA data.

Program.Data also accepts non-differentiable tensor inputs with their own element type, such as token indices. map applies an operation over a specified prefix shape; linear, convolution, pooling, normalization and attention accept those prefixes without separate batched variants.

The operation-polymorphic surface shared by the eager tape and the typed-graph recorder.

Prefix-polymorphic derived operations

def Runtime.Autograd.Model.map {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (leadingShape : Spec.Shape) {s t : Spec.Shape} (x : Ref (leadingShape.concat s)) (f : Ref s → m (Ref t)) : m (Ref (leadingShape.concat t))

Apply a single-sample operation independently over an arbitrary prefix shape.

def Runtime.Autograd.Model.linear {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {inDim outDim : ℕ} (weight : Ref [outDim, inDim]) (bias : Ref [outDim]) (input : Ref (leadingShape.concat [inDim])) : m (Ref (leadingShape.concat [outDim]))

Affine transformation of the final axis, independently over any prefix shape.

def Runtime.Autograd.Model.conv {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (weight : Ref (Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ)))) (bias : Ref [outC]) (input : Ref (leadingShape.concat (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ))))) : m (Ref (leadingShape.concat (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).to (List ℕ)))))

Rank-polymorphic convolution over channels-first inputs with any prefix shape.

def Runtime.Autograd.Model.convTranspose {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (weight : Ref (Spec.Shape.ofList (inC :: outC :: kernel.to (List ℕ)))) (bias : Ref [outC]) (input : Ref (leadingShape.concat (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ))))) : m (Ref (leadingShape.concat (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).to (List ℕ)))))

Rank-polymorphic transpose convolution over channels-first inputs with any prefix shape.

def Runtime.Autograd.Model.maxPool {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {d channels : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} (input : Ref (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) : m (Ref (leadingShape.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ)))))

Rank-polymorphic max pooling over channels-first inputs with any prefix shape.

def Runtime.Autograd.Model.avgPool {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {d channels : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} (input : Ref (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) : m (Ref (leadingShape.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ)))))

Rank-polymorphic average pooling over channels-first inputs with any prefix shape.

def Runtime.Autograd.Model.smoothMaxPool {α : Type} [TorchLean.Storage α] [Context α] [DecidableEq α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {d channels : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} (input : Ref (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) (temp : α) : m (Ref (leadingShape.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ)))))

Batched rank-polymorphic smooth maximum pooling over channels-first inputs.

def Runtime.Autograd.Model.layerNorm {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leading : Spec.Shape} {width : ℕ} (hWidth : width > 0) (x : Ref (leading.appendDim width)) (gamma beta : Ref [width]) (epsilon : α := TorchLean.normalizationEpsilon) : m (Ref (leading.appendDim width))

Layer normalization over the final axis of an arbitrary tensor.

Backends expose a matrix-shaped fused primitive. This wrapper gives that primitive its rank-polymorphic TorchLean interface by flattening the leading axes once and restoring them after the operation. Empty leading axes produce the unique empty tensor without invoking a backend kernel that requires a positive row count.

def Runtime.Autograd.Model.attention {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {n numHeads dModel headDim : ℕ} (hN : n ≠ 0) (wq wk wv : Ref [dModel, numHeads * headDim]) (wo : Ref [numHeads * headDim, dModel]) (x : Ref (leadingShape.concat [n, dModel])) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) (outputBias : Option (Ref [dModel]) := none) : m (Ref (leadingShape.concat [n, dModel]))

Self-attention over any prefix shape, with an optional output-feature bias.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Program (α : Type) [TorchLean.Storage α] [Context α] (ss : List Spec.Shape) (τ : Spec.Shape) : Type 1

An execution-polymorphic differentiable tensor program.

@[reducible, inline]

abbrev Runtime.Autograd.Model.Program.Data (α β : Type) [TorchLean.Storage α] [TorchLean.Storage β] [Context α] (ss dataSs : List Spec.Shape) (τ : Spec.Shape) : Type 1

An execution-polymorphic program with differentiable tensors followed by non-differentiable data tensors. The data element type is explicit; for example, token models may use Fin vocab so an out-of-range token is unrepresentable.