Eager Tensor Operations

PyTorch-style tensor operations backed by the eager CPU/CUDA tapes. These wrappers record runtime nodes, dispatch CUDA kernels when requested, and preserve the typed TensorRef surface.

Tensor ops (eager tape wrappers)

The following definitions are the eager front-end for Runtime.Autograd.Tape.* primitives. Each one:

• reads the current tape from s.tape,
• appends a new node/leaf via a Tape.* constructor,
• writes the updated tape back, and
• returns a fresh TensorRef pointing to the new node id.

PyTorch comparison: this is the standard eager autograd mechanism (a dynamic tape of ops).

def Runtime.Autograd.Torch.Internal.EagerSession.dispatchCudaOpt {α β : Type} (s : EagerSession α) (opName : String) (cpu : IO β) (cuda : IO (Option β)) : IO β

Dispatch an eager op with optional CUDA support.

When Options.device = .cuda, any op whose CUDA implementation returns none will throw.

TorchLean's CUDA eager mode has no per-op CPU fallback: either the op is supported by CUDA, or it errors immediately.

def Runtime.Autograd.Torch.Internal.EagerSession.add {α : Type} (s : EagerSession α) [Add α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise addition a + b. PyTorch: torch.add.

def Runtime.Autograd.Torch.Internal.EagerSession.sub {α : Type} (s : EagerSession α) [Sub α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise subtraction a - b. PyTorch: torch.sub.

def Runtime.Autograd.Torch.Internal.EagerSession.mul {α : Type} (s : EagerSession α) [Mul α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise multiplication a * b. PyTorch: torch.mul.

def Runtime.Autograd.Torch.Internal.EagerSession.scale {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Mul α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (c : α) : IO (TensorRef α sh)

Record scaling by a scalar constant. PyTorch: x * c.

def Runtime.Autograd.Torch.Internal.EagerSession.abs {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise absolute value. PyTorch: torch.abs.

def Runtime.Autograd.Torch.Internal.EagerSession.sqrt {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise square root. PyTorch: torch.sqrt.

def Runtime.Autograd.Torch.Internal.EagerSession.clamp {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (minVal maxVal : α) : IO (TensorRef α sh)

Record elementwise clamp to [minVal,maxVal]. PyTorch: torch.clamp.

def Runtime.Autograd.Torch.Internal.EagerSession.max {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise maximum. PyTorch: torch.maximum.

def Runtime.Autograd.Torch.Internal.EagerSession.min {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise minimum. PyTorch: torch.minimum.

def Runtime.Autograd.Torch.Internal.EagerSession.relu {α : Type} (s : EagerSession α) [Mul α] [Zero α] [Max α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise ReLU. PyTorch: torch.relu / torch.nn.functional.relu.

def Runtime.Autograd.Torch.Internal.EagerSession.sigmoid {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise sigmoid. PyTorch: torch.sigmoid.

def Runtime.Autograd.Torch.Internal.EagerSession.tanh {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise tanh. PyTorch: torch.tanh.

def Runtime.Autograd.Torch.Internal.EagerSession.softmax {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record softmax (shape-preserving).

PyTorch comparison: torch.softmax(x, dim=...) (dimension convention is chosen by the underlying tape op).

def Runtime.Autograd.Torch.Internal.EagerSession.logSoftmax {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record stable log-softmax (shape-preserving, last-axis convention).

PyTorch comparison: torch.nn.functional.log_softmax(x, dim=-1).

def Runtime.Autograd.Torch.Internal.EagerSession.softplus {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise softplus. PyTorch: torch.nn.functional.softplus.

def Runtime.Autograd.Torch.Internal.EagerSession.exp {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise exponential. PyTorch: torch.exp.

def Runtime.Autograd.Torch.Internal.EagerSession.log {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise log. PyTorch: torch.log.

def Runtime.Autograd.Torch.Internal.EagerSession.inv {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise inverse 1/x. PyTorch: torch.reciprocal.

def Runtime.Autograd.Torch.Internal.EagerSession.safeLog {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (ε : α := Numbers.epsilon) : IO (TensorRef α sh)

Record elementwise log with epsilon guard.

PyTorch comparison: torch.log(torch.clamp(x, min=ε)).

def Runtime.Autograd.Torch.Internal.EagerSession.sum {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α Spec.Shape.scalar)

Sum-reduce all elements to a scalar. PyTorch: x.sum().

def Runtime.Autograd.Torch.Internal.EagerSession.flatten {α : Type} (s : EagerSession α) [Inhabited α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α (Spec.Shape.dim sh.size Spec.Shape.scalar))

Flatten a tensor to a 1D vector. PyTorch: torch.flatten.

def Runtime.Autograd.Torch.Internal.EagerSession.reshape {α : Type} (s : EagerSession α) [Inhabited α] [DecidableEq Spec.Shape] {sh1 sh2 : Spec.Shape} (x : TensorRef α sh1) (h : sh1.size = sh2.size) : IO (TensorRef α sh2)

Reshape a tensor while preserving total number of elements.

PyTorch comparison: torch.reshape / view (when valid).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose2d {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {m n : ℕ} (x : TensorRef α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : IO (TensorRef α (Spec.Shape.dim n (Spec.Shape.dim m Spec.Shape.scalar)))

Transpose a 2D matrix. PyTorch: x.t() / x.transpose(0,1).

def Runtime.Autograd.Torch.Internal.EagerSession.swapAdjacentAtDepth {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {sh : Spec.Shape} (depth : ℕ) (x : TensorRef α sh) : IO (TensorRef α (sh.swapAdjacentAtDepth depth))

Swap two adjacent axes at a given depth. PyTorch analogue: x.transpose(dim, dim+1).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dFirstToLast {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim a Spec.Shape.scalar))))

Permute a 3D tensor (a,b,c) → (b,c,a). PyTorch: x.permute(1,2,0).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dLastToFirst {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim c (Spec.Shape.dim a (Spec.Shape.dim b Spec.Shape.scalar))))

Permute a 3D tensor (a,b,c) → (c,a,b). PyTorch: x.permute(2,0,1).

def Runtime.Autograd.Torch.Internal.EagerSession.transpose3dLastTwo {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {a b c : ℕ} (x : TensorRef α (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim a (Spec.Shape.dim c (Spec.Shape.dim b Spec.Shape.scalar))))

Swap the last two axes of a 3D tensor (a,b,c) → (a,c,b). PyTorch: x.transpose(1,2).

def Runtime.Autograd.Torch.Internal.EagerSession.broadcastTo {α : Type} (s : EagerSession α) [Inhabited α] [Add α] [Zero α] [DecidableEq Spec.Shape] {sh1 sh2 : Spec.Shape} (cb : sh1.CanBroadcastTo sh2) (x : TensorRef α sh1) : IO (TensorRef α sh2)

Broadcast a tensor to a larger shape. PyTorch: implicit broadcasting / expand.

def Runtime.Autograd.Torch.Internal.EagerSession.reduceSum {α : Type} (s : EagerSession α) [Add α] [Zero α] [Inhabited α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis sh] [wf : sh.WellFormed] (x : TensorRef α sh) : IO (TensorRef α (Spec.Tensor.shapeAfterSum sh axis))

Sum-reduce along axis. PyTorch: torch.sum(x, dim=axis).

def Runtime.Autograd.Torch.Internal.EagerSession.reduceMean {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis sh] [wf : sh.WellFormed] (x : TensorRef α sh) : IO (TensorRef α (Spec.Tensor.shapeAfterSum sh axis))

Mean-reduce along axis. PyTorch: torch.mean(x, dim=axis).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalar {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) : IO (TensorRef α Spec.Shape.scalar)

Gather a scalar from a 1D vector with a Fin n index. PyTorch: x[i].

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRow {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : Fin rows) : IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

Gather a row from a 2D tensor with a Fin rows index. PyTorch: x[i] for 2D tensors.

def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalarNat {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : ℕ) : IO (TensorRef α Spec.Shape.scalar)

Gather a scalar from a 1D vector with a raw Nat index (totalized by the tape op).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherScalarRef {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (i : NatRef) : IO (TensorRef α Spec.Shape.scalar)

Dynamic gather scalar using an index stored in NatRef.

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowRef {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : NatRef) : IO (TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar))

Dynamic gather row using an index stored in NatRef (out-of-range gives a zero row).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

Gather k scalars using an explicit index tensor. PyTorch analogue: gather / advanced indexing.

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsNat {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Gather k rows using an explicit index tensor. PyTorch: index_select(dim=0, index=...).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherVecRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n k : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (idx : NatVecRef k) : IO (TensorRef α (Spec.Shape.dim k Spec.Shape.scalar))

Gather k scalars using indices stored in the nat-environment (NatVecRef).

def Runtime.Autograd.Torch.Internal.EagerSession.gatherRowsRef {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols k : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : NatVecRef k) : IO (TensorRef α (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Gather k rows using indices stored in the nat-environment (NatVecRef).

def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddVec {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {n : ℕ} (x : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (v : TensorRef α Spec.Shape.scalar) (i : Fin n) : IO (TensorRef α (Spec.Shape.dim n Spec.Shape.scalar))

Scatter-add into a vector: return a copy of x with x[i] += v.

def Runtime.Autograd.Torch.Internal.EagerSession.scatterAddRow {α : Type} (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {rows cols : ℕ} (x : TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (v : TensorRef α (Spec.Shape.dim cols Spec.Shape.scalar)) (i : Fin rows) : IO (TensorRef α (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))

Scatter-add into a matrix row: return a copy of x with x[i,:] += v.

def Runtime.Autograd.Torch.Internal.EagerSession.linear {α : Type} (s : EagerSession α) [Inhabited α] [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] [FastKernels.FastMatmul α] {inDim outDim : ℕ} (w : TensorRef α (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b : TensorRef α (Spec.Shape.dim outDim Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.dim inDim Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim outDim Spec.Shape.scalar))

Fully-connected linear layer y = w x + b (matvec).

If opts.fastKernels is enabled, uses a runtime-only fast kernel implementation. PyTorch comparison: torch.nn.functional.linear(x, weight=w, bias=b).

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

Mean-squared-error loss returning a scalar.

If opts.fastKernels is enabled, uses a runtime-only fast kernel implementation. PyTorch comparison: torch.nn.functional.mse_loss(..., reduction=\"mean\").

def Runtime.Autograd.Torch.Internal.EagerSession.layerNorm {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : TensorRef α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) (gamma beta : TensorRef α (Spec.Shape.dim embedDim Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

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

def Runtime.Autograd.Torch.Internal.EagerSession.batchnormChannelFirst {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {channels height width : ℕ} (h_c : channels > 0) (h_h : height > 0) (h_w : width > 0) (x : TensorRef α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar)))) (gamma beta : TensorRef α (Spec.Shape.dim channels Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim channels (Spec.Shape.dim height (Spec.Shape.dim width Spec.Shape.scalar))))

BatchNorm for channel-first images (C,H,W) (no batch axis). PyTorch: nn.BatchNorm2d (conceptually).

def Runtime.Autograd.Torch.Internal.EagerSession.multiHeadAttention {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar))) (wo : TensorRef α (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar))) (x : TensorRef α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : IO (TensorRef α (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))

Multi-head self-attention (typed, proof-friendly). PyTorch: nn.MultiheadAttention (conceptually).

def Runtime.Autograd.Torch.Internal.EagerSession.conv {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (w : TensorRef α (Spec.Shape.ofList (outC :: inC :: kernel.toList))) (b : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.ofList (inC :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).toList)))

N-D convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv2d. PyTorch comparison: torch.nn.functional.conv{d}d specialized to a single sample.

def Runtime.Autograd.Torch.Internal.EagerSession.conv2d {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim inC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (input : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D convolution for channel-first images (C,H,W) (no batch axis). PyTorch: torch.nn.functional.conv2d.

def Runtime.Autograd.Torch.Internal.EagerSession.convTranspose {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} {hInC : inC ≠ 0} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (w : TensorRef α (Spec.Shape.ofList (inC :: outC :: kernel.toList))) (b : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (x : TensorRef α (Spec.Shape.ofList (inC :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).toList)))

N-D transpose convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv_transpose2d. PyTorch comparison: torch.nn.functional.conv_transpose{d}d specialized to a single sample.

def Runtime.Autograd.Torch.Internal.EagerSession.convTranspose2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (kernel : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim outC (Spec.Shape.dim kH (Spec.Shape.dim kW Spec.Shape.scalar))))) (bias : TensorRef α (Spec.Shape.dim outC Spec.Shape.scalar)) (input : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim outC (Spec.Shape.dim ((inH - 1) * stride - 2 * padding + kH) (Spec.Shape.dim ((inW - 1) * stride - 2 * padding + kW) Spec.Shape.scalar))))

2D transpose convolution for channel-first images (C,H,W) (no batch axis). PyTorch: torch.nn.functional.conv_transpose2d.

def Runtime.Autograd.Torch.Internal.EagerSession.matmul {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {m n p : ℕ} (a : TensorRef α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : TensorRef α (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : IO (TensorRef α (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar)))

2D matrix multiplication. PyTorch: torch.matmul for 2D tensors.

def Runtime.Autograd.Torch.Internal.EagerSession.bmm {α : Type} (s : EagerSession α) [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] {batch m n p : ℕ} (a : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)))) (b : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))))

Batched matrix multiplication. PyTorch: torch.bmm.

def Runtime.Autograd.Torch.Internal.EagerSession.concatVectors {α : Type} (s : EagerSession α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {n m : ℕ} (a : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (b : TensorRef α (Spec.Shape.dim m Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim (n + m) Spec.Shape.scalar))

Concatenate two vectors along dim 0. PyTorch: torch.cat([a,b], dim=0).

def Runtime.Autograd.Torch.Internal.EagerSession.concatDim0 {α : Type} (s : EagerSession α) [DecidableEq Spec.Shape] {n m : ℕ} {sh : Spec.Shape} (a : TensorRef α (Spec.Shape.dim n sh)) (b : TensorRef α (Spec.Shape.dim m sh)) : IO (TensorRef α (Spec.Shape.dim (n + m) sh))

Concatenate along dim 0 for tensors with leading dimension. PyTorch: torch.cat(..., dim=0).

def Runtime.Autograd.Torch.Internal.EagerSession.sliceRange0 {α : Type} (s : EagerSession α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} {sh : Spec.Shape} (x : TensorRef α (Spec.Shape.dim n sh)) (start len : ℕ) (h : len + start ≤ n) : IO (TensorRef α (Spec.Shape.dim len sh))

Slice along dim 0: x[start:start+len]. PyTorch: standard slicing.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D max pooling for channels-first tensors (C, spatial...) (no batch axis).

PyTorch comparison: torch.nn.functional.max_pool1d / max_pool2d / max_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D average pooling for channels-first tensors (C, spatial...) (no batch axis).

PyTorch comparison: torch.nn.functional.avg_pool1d / avg_pool2d / avg_pool3d depending on the spatial rank d.

def Runtime.Autograd.Torch.Internal.EagerSession.smoothMaxPool {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : TensorRef α (Spec.Shape.ofList (C :: inSpatial.toList))) (beta : α) : IO (TensorRef α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D smooth max pooling (log-sum-exp surrogate) for channels-first tensors (C, spatial...).

This is a differentiable approximation to max pooling; PyTorch does not expose it as a single primitive, but it can be emulated with logsumexp over local windows.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling (no batch axis). PyTorch: torch.nn.functional.max_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.maxPool2dPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling with padding (no batch axis). PyTorch: max_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.maxPoolPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for max_pool2d_pad (PyTorch-style shorthand).

def Runtime.Autograd.Torch.Internal.EagerSession.smoothMaxPool2d {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (beta : α) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Smooth max-pooling (softmax pooling). Not a standard PyTorch primitive; see Torch.LinkedSession.smooth_max_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool2d {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

2D average-pooling (no batch axis). PyTorch: torch.nn.functional.avg_pool2d.

def Runtime.Autograd.Torch.Internal.EagerSession.avgPool2dPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D average-pooling with padding (no batch axis). PyTorch: avg_pool2d(..., padding=...).

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerSession.avgPoolPad {α : Type} (s : EagerSession α) [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Alias for avg_pool2d_pad (PyTorch-style shorthand).