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 interpreter records reusable shape-indexed SSA data.

The execution-polymorphic surface shared by eager and typed-graph execution.

Operation-reference notation

@[reducible, inline]

abbrev Runtime.Autograd.Model.RefTy (m : Type → Type) (α : Type) [TorchLean.Storage α] [Context α] [Ops m α] (s : Spec.Shape) : Type

A tensor reference under the currently selected operation interpreter.

def Runtime.Autograd.Model.LeadingAxis.Internal.mapOuterAxis {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch : ℕ} {s t : Spec.Shape} (x : RefTy m α (s.prependDim batch)) (f : RefTy m α s → m (RefTy m α t)) : m (RefTy m α (t.prependDim batch))

Apply a single-sample operation independently along a leading axis.

Prefix-polymorphic derived operations

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

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

def Runtime.Autograd.Model.linearEach {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {inDim outDim : ℕ} (weight : RefTy m α [outDim, inDim]) (bias : RefTy m α [outDim]) (input : RefTy m α (leadingShape.concat [inDim])) : m (RefTy m α (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 : RefTy m α (Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ)))) (bias : RefTy m α [outC]) (input : RefTy m α (leadingShape.concat (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ))))) : m (RefTy m α (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 : RefTy m α (Spec.Shape.ofList (inC :: outC :: kernel.to (List ℕ)))) (bias : RefTy m α [outC]) (input : RefTy m α (leadingShape.concat (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ))))) : m (RefTy m α (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 : RefTy m α (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) : m (RefTy m α (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 : RefTy m α (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) : m (RefTy m α (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 : RefTy m α (leadingShape.concat (Spec.Shape.ofList (channels :: inSpatial.to (List ℕ))))) (temp : α) : m (RefTy m α (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 : RefTy m α (leading.appendDim width)) (gamma beta : RefTy m α [width]) (epsilon : α := TorchLean.normalizationEpsilon) : m (RefTy m α (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.multiHeadAttention {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {leadingShape : Spec.Shape} {n numHeads dModel headDim : ℕ} (hN : n ≠ 0) (wq wk wv : RefTy m α [dModel, numHeads * headDim]) (wo : RefTy m α [numHeads * headDim, dModel]) (x : RefTy m α (leadingShape.concat [n, dModel])) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) : m (RefTy m α (leadingShape.concat [n, dModel]))

Multi-head self-attention over any prefix shape.

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

Multi-head attention followed by a trainable 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.ProgramWithDataInputs (α β : 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.