Eager Tensor Operations

PyTorch-style tensor operations recorded on the session's tape. Each one runs the CPU or the CUDA kernel of the operation, as the session's device selects, and preserves the typed TensorRef surface.

Neural-network layers

def Runtime.Autograd.Torch.Internal.EagerSession.softmaxMatmul {α : Type} [TorchLean.Storage α] [Context α] (session : EagerSession α) {batch : Spec.Shape} {rows keys width : ℕ} (scores : TensorRef α (batch.concat [rows, keys])) (value : TensorRef α (batch.concat [keys, width])) (multiplier : TensorRef α (batch.concat [rows, keys])) (mask : Option (TorchLean.Tensor Bool [rows, keys]) := none) : IO (TensorRef α (batch.concat [rows, width]))

Normalize score rows and aggregate values, retaining the compound local pullback.

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

Fully-connected linear layer y = w x + b. PyTorch: torch.nn.functional.linear.

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

Mean-squared-error loss returning a scalar. PyTorch: torch.nn.functional.mse_loss.

def Runtime.Autograd.Torch.Internal.EagerSession.crossEntropy {α : Type} [TorchLean.Storage α] [Context α] (session : EagerSession α) {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (logits target : TensorRef α s) (mean : Bool := true) : IO (TensorRef α Spec.Shape.scalar)

Mean weighted logits cross-entropy, recorded on the selected backend's existing tape.

def Runtime.Autograd.Torch.Internal.EagerSession.distill {α : Type} [TorchLean.Storage α] [Context α] (session : EagerSession α) {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (student teacher : TensorRef α s) : IO (TensorRef α Spec.Shape.scalar)

Mean teacher-logit cross-entropy, recorded on the selected backend's existing tape.

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

Layer normalization over embedding dimension. PyTorch: nn.LayerNorm / functional.layer_norm.

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

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

def Runtime.Autograd.Torch.Internal.EagerSession.attentionCpu {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Context α] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α [dModel, numHeads * headDim]) (wo : TensorRef α [numHeads * headDim, dModel]) (x : TensorRef α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n])) : IO (TensorRef α [n, dModel])

CPU attention for one sequence, recorded on TorchLean's reference tape.

def Runtime.Autograd.Torch.Internal.EagerSession.attention {α : Type} [TorchLean.Storage α] (s : EagerSession α) [Context α] [TensorTransfer α] {n numHeads dModel headDim : ℕ} {batch : Option ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α [dModel, numHeads * headDim]) (wo : TensorRef α [numHeads * headDim, dModel]) (x : TensorRef α (match batch with | none => [n, dModel] | some b => [b, n, dModel])) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) (hBatch : batch.getD 1 ≠ 0 := by decide) : IO (TensorRef α (match batch with | none => [n, dModel] | some b => [b, n, dModel]))

Self-attention with an optional leading batch dimension.

The head count is numHeads; batch := some b selects inputs of shape [b, n, dModel]. CPU execution maps the reference operation over samples; GPU execution records one tape node for the batch and evaluates it through multiple LibTorch primitives.