Session Neural-Network Operations

This file contains higher-level neural-network session calls such as linear layers, normalization, attention, and convolutional blocks. The operations share the same session dispatch discipline as the elementary ops while preserving PyTorch-style call sites.

def Runtime.Autograd.Model.Session.linear {α : Type} [TorchLean.Storage α] (s : Session α) [Inhabited α] [Add α] [Mul α] [Zero α] {inDim outDim : ℕ} (w : Torch.TensorRef α [outDim, inDim]) (b : Torch.TensorRef α [outDim]) (x : Torch.TensorRef α [inDim]) : IO (Torch.TensorRef α [outDim])

Fully-connected (affine) layer on vectors: $y=w\mathbin{\cdot}x+b$.

PyTorch analogue: torch.nn.functional.linear (weight shape (outDim, inDim)).

def Runtime.Autograd.Model.Session.mseLoss {α : Type} [TorchLean.Storage α] (s : Session α) [Inhabited α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [One α] [NatCast α] [Torch.TensorTransfer α] {sh : Spec.Shape} (yhat target : Torch.TensorRef α sh) : IO (Torch.TensorRef α Spec.Shape.scalar)

Mean squared error loss returning a scalar.

PyTorch analogue: torch.nn.functional.mse_loss(..., reduction='mean').

def Runtime.Autograd.Model.Session.layerNorm {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [Torch.TensorTransfer α] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : Torch.TensorRef α [seqLen, embedDim]) (gamma beta : Torch.TensorRef α [embedDim]) (epsilon : α := TorchLean.normalizationEpsilon) : IO (Torch.TensorRef α [seqLen, embedDim])

LayerNorm over a seqLen × embedDim tensor.

PyTorch analogue: torch.nn.LayerNorm(embedDim) applied per token.

def Runtime.Autograd.Model.Session.batchNorm {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [Torch.TensorTransfer α] {channels : ℕ} {sSpatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim channels sSpatial).wellFormed) (x : Torch.TensorRef α (Spec.Shape.dim channels sSpatial)) (gamma beta : Torch.TensorRef α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) : IO (Torch.TensorRef α (Spec.Shape.dim channels sSpatial))

Batch normalization over every spatial axis of a channel-first tensor.

def Runtime.Autograd.Model.Session.conv {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (w : Torch.TensorRef α (Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ)))) (b : Torch.TensorRef α [outC]) (x : Torch.TensorRef α (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ)))) : IO (Torch.TensorRef α (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).to (List ℕ))))

N-D convolution over a channels-first tensor (inC, spatial...).

PyTorch analogue: torch.nn.functional.conv{d}d specialized to a single sample.

def Runtime.Autograd.Model.Session.convTranspose {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (w : Torch.TensorRef α (Spec.Shape.ofList (inC :: outC :: kernel.to (List ℕ)))) (b : Torch.TensorRef α [outC]) (x : Torch.TensorRef α (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ)))) : IO (Torch.TensorRef α (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).to (List ℕ))))

N-D transpose convolution over a channels-first tensor (inC, spatial...).

PyTorch analogue: torch.nn.functional.conv_transpose{d}d specialized to a single sample.

def Runtime.Autograd.Model.Session.multiHeadAttention {α : Type} [TorchLean.Storage α] (s : Session α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : Torch.TensorRef α [dModel, numHeads * headDim]) (wo : Torch.TensorRef α [numHeads * headDim, dModel]) (x : Torch.TensorRef α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) : IO (Torch.TensorRef α [n, dModel])

Multi-head self-attention (single sequence, single batch).

This is a convenience op used by the transformer examples; it corresponds approximately to the forward pass of torch.nn.MultiheadAttention in "self-attention" mode.