Small helpers

def Runtime.Autograd.Cuda.Tape.u32 (n : ℕ) : Result UInt32

Checked Nat → UInt32 conversion for CUDA boundaries. Errors if n ≥ 2^32.

def Runtime.Autograd.Cuda.Tape.numelU32 (s : Spec.Shape) : Result UInt32

Number of elements in a runtime shape, checked for the UInt32 CUDA ABI.

def Runtime.Autograd.Cuda.Tape.foldRowsColsLastAxis (s : Spec.Shape) : Result (UInt32 × UInt32)

Fold all leading axes into a row count and keep the last axis as the column count.

CUDA softmax/log-softmax kernels are 2D row kernels. This helper gives the shared convention used for vectors, matrices, and higher-rank tensors: softmax is always along the last axis.

def Runtime.Autograd.Cuda.Tape.broadcastScalarToShape (g : Buffer) (outShape : Spec.Shape) : Buffer

Broadcast a scalar CUDA buffer to outShape. Used by scalar reductions during backprop.

def Runtime.Autograd.Cuda.Tape.sigmoidBuf (x : Buffer) (n : UInt32) : Buffer

Logistic sigmoid implemented from primitive CUDA elementwise ops.

def Runtime.Autograd.Cuda.Tape.tanhBuf (x : Buffer) (n : UInt32) : Buffer

Hyperbolic tangent implemented as (exp(2x)-1)/(exp(2x)+1).

def Runtime.Autograd.Cuda.Tape.softplusBuf (x : Buffer) (n : UInt32) : Buffer

Numerically stable softplus: max(x,0) + log(1 + exp(-abs(x))).

def Runtime.Autograd.Cuda.Tape.softmax2dFwdWithCleanup (x : Buffer) (rows cols : UInt32) : Buffer × List Buffer

Row-wise stable softmax plus scratch buffers that should be released after the backward pass.

The output is the first component. The scratch list records temporary buffers created while computing the stable formula so long training runs do not wait for native finalizers.

def Runtime.Autograd.Cuda.Tape.softmax2dFwd (x : Buffer) (rows cols : UInt32) : Buffer

def Runtime.Autograd.Cuda.Tape.softmax2dBwd (y dLdy : Buffer) (rows cols : UInt32) : Buffer

Row-wise softmax VJP: dX = y * (dY - sum(dY*y, axis=1)).

def Runtime.Autograd.Cuda.Tape.logSoftmax2dFwdWithCleanup (x : Buffer) (rows cols : UInt32) : Buffer × List Buffer

Row-wise stable log-softmax plus scratch buffers that should be released after backprop.

This computes x - rowMax - log(sum(exp(x-rowMax))) directly, avoiding the less stable log(softmax(x)) route.

def Runtime.Autograd.Cuda.Tape.logSoftmax2dFwd (x : Buffer) (rows cols : UInt32) : Buffer

def Runtime.Autograd.Cuda.Tape.logSoftmax2dBwd (y dLdy : Buffer) (rows cols : UInt32) : Buffer

Row-wise log-softmax VJP: dX = dY - exp(y) * sum(dY, axis=1).

Elementwise ops

def Runtime.Autograd.Cuda.Tape.add {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise addition.

def Runtime.Autograd.Cuda.Tape.sub {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise subtraction.

def Runtime.Autograd.Cuda.Tape.mul {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise multiplication.

def Runtime.Autograd.Cuda.Tape.scale {s : Spec.Shape} (t : Tape) (xId : ℕ) (c : Float) : Result (Tape × ℕ)

Multiply by a scalar constant.

def Runtime.Autograd.Cuda.Tape.abs {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise abs.

def Runtime.Autograd.Cuda.Tape.sqrt {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise sqrt.

def Runtime.Autograd.Cuda.Tape.clamp {s : Spec.Shape} (t : Tape) (xId : ℕ) (lo hi : Float) : Result (Tape × ℕ)

Clamp each element to [lo, hi].

def Runtime.Autograd.Cuda.Tape.max {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise max.

def Runtime.Autograd.Cuda.Tape.min {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise min.

def Runtime.Autograd.Cuda.Tape.div {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Elementwise division.

def Runtime.Autograd.Cuda.Tape.relu {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise ReLU.

def Runtime.Autograd.Cuda.Tape.exp {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise exp.

def Runtime.Autograd.Cuda.Tape.log {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise log.

def Runtime.Autograd.Cuda.Tape.inv {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise reciprocal 1/x.

def Runtime.Autograd.Cuda.Tape.safeLog {s : Spec.Shape} (t : Tape) (xId : ℕ) (ε : Float) : Result (Tape × ℕ)

Elementwise "safe log" that protects against log(0) by adding a small ε internally.

Spec semantics: log(softplus(x) + ε).

def Runtime.Autograd.Cuda.Tape.sigmoid {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise sigmoid (logistic).

def Runtime.Autograd.Cuda.Tape.tanh {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise tanh.

def Runtime.Autograd.Cuda.Tape.softplus {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Elementwise softplus.

Reductions / views

def Runtime.Autograd.Cuda.Tape.sum {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Reduce-sum of all entries, producing a scalar.

def Runtime.Autograd.Cuda.Tape.flatten {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Flatten s into a 1D vector of length Shape.size s.

def Runtime.Autograd.Cuda.Tape.reshape {s₁ s₂ : Spec.Shape} (t : Tape) (xId : ℕ) (_h : s₁.size = s₂.size) : Result (Tape × ℕ)

Reshape a buffer while preserving number of elements.

This is a no-copy view operation: it reuses the same contiguous buffer.

def Runtime.Autograd.Cuda.Tape.transpose2d {m n : ℕ} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Transpose a 2D buffer.

def Runtime.Autograd.Cuda.Tape.swapAdjacentAtDepth {s : Spec.Shape} (t : Tape) (depth xId : ℕ) : Result (Tape × ℕ)

Swap adjacent axes at a given depth in an N-D buffer.

If depth is out of range, this is treated as the identity (matches the spec-layer helper).

def Runtime.Autograd.Cuda.Tape.transpose3dFirstToLast {a b c : ℕ} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Permute a 3D tensor (a,b,c) → (b,c,a).

def Runtime.Autograd.Cuda.Tape.transpose3dLastToFirst {a b c : ℕ} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Permute a 3D tensor (a,b,c) → (c,a,b).

def Runtime.Autograd.Cuda.Tape.transpose3dLastTwo {a b c : ℕ} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Swap the last two axes of a 3D tensor (a,b,c) → (a,c,b).

def Runtime.Autograd.Cuda.Tape.broadcastTo {s₁ s₂ : Spec.Shape} (t : Tape) (cb : s₁.CanBroadcastTo s₂) (xId : ℕ) : Result (Tape × ℕ)

Broadcast x : s₁ to s₂.

Forward: broadcastTo. Backward: sum-reduce broadcasted axes (reduceFromBroadcastTo).

def Runtime.Autograd.Cuda.Tape.reduceSum {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis s] [wf : s.WellFormed] (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Reduce-sum along axis.

def Runtime.Autograd.Cuda.Tape.reduceMean {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis s] [wf : s.WellFormed] (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Reduce-mean along axis.

Linear algebra

def Runtime.Autograd.Cuda.Tape.matmul {m n p : ℕ} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

2D matrix multiply.

def Runtime.Autograd.Cuda.Tape.bmm {batch m n p : ℕ} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Batched matrix multiply.

def Runtime.Autograd.Cuda.Tape.spectralConv1dRfft {grid width modes : ℕ} (t : Tape) (xId wReId wImId : ℕ) : Result (Tape × ℕ)

Fused real-FFT spectral convolution used by the CUDA FNO1D path.

Shapes:

• x : (grid, width),
• wRe, wIm : (modes, width, width),
• output y : (grid, width).

The low-level buffer primitive owns the numerical contract and VJP: rfft(x) is unnormalized, the inverse is normalized, and the backward kernels include the half-spectrum adjoint factors for real FFTs. This tape node simply records those three parent dependencies and checks the runtime shapes before calling the native kernels.

Linear layer / losses

def Runtime.Autograd.Cuda.Tape.linear {outDim inDim : ℕ} (t : Tape) (wId bId xId : ℕ) : Result (Tape × ℕ)

Linear layer: y = W·x + b with W : (outDim,inDim), x : inDim, b : outDim.

def Runtime.Autograd.Cuda.Tape.mseLoss {s : Spec.Shape} (t : Tape) (yhatId targetId : ℕ) : Result (Tape × ℕ)

Mean-squared-error loss with "mean" reduction (single scalar output).

Concat / slice (1D)

def Runtime.Autograd.Cuda.Tape.concat1d {n m : ℕ} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Concatenate two 1D buffers.

def Runtime.Autograd.Cuda.Tape.concatVectors {n m : ℕ} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Concatenate two 1D tensors (CPU tape name).

def Runtime.Autograd.Cuda.Tape.slice1d {n start len : ℕ} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Slice a 1D buffer.

Concat / slice along dim 0

def Runtime.Autograd.Cuda.Tape.concatDim0 {n m : ℕ} {s : Spec.Shape} (t : Tape) (aId bId : ℕ) : Result (Tape × ℕ)

Concatenate along dim 0 for tensors with leading dimension (CPU tape name).

def Runtime.Autograd.Cuda.Tape.sliceRange0 {n : ℕ} {s : Spec.Shape} (t : Tape) (xId start len : ℕ) (_h : len + start ≤ n) : Result (Tape × ℕ)

Slice along dim 0: x[start:start+len] (CPU tape name).

Gather / scatter (host Nat indices)

Indices are non-differentiable and remain on the host. Kernels totalize out-of-bounds indices as documented in NN.Runtime.Autograd.Engine.Cuda.Kernels.

def Runtime.Autograd.Cuda.Tape.gatherScalar {n : ℕ} (t : Tape) (xId : ℕ) (i : Fin n) : Result (Tape × ℕ)

Gather a scalar from a 1D vector using a compile-time index.

def Runtime.Autograd.Cuda.Tape.gatherRow {rows cols : ℕ} (t : Tape) (xId : ℕ) (i : Fin rows) : Result (Tape × ℕ)

Gather a row from a 2D matrix using a compile-time index.

def Runtime.Autograd.Cuda.Tape.gatherScalarNat {n : ℕ} (t : Tape) (xId i : ℕ) : Result (Tape × ℕ)

Gather a scalar from a 1D vector using a runtime Nat index (totalized by the kernel).

def Runtime.Autograd.Cuda.Tape.natTensorToIndexArray {k : ℕ} (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : Array ℕ

def Runtime.Autograd.Cuda.Tape.gatherVecNat {n k : ℕ} (t : Tape) (xId : ℕ) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : Result (Tape × ℕ)

Gather k scalars from a length-n vector.

def Runtime.Autograd.Cuda.Tape.gatherRowsNat {rows cols k : ℕ} (t : Tape) (xId : ℕ) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : Result (Tape × ℕ)

Gather k rows from a (rows, cols) matrix (row-major).

def Runtime.Autograd.Cuda.Tape.scatterAddVec {n : ℕ} (t : Tape) (xId vId : ℕ) (i : Fin n) : Result (Tape × ℕ)

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

def Runtime.Autograd.Cuda.Tape.scatterAddRow {rows cols : ℕ} (t : Tape) (xId vId : ℕ) (i : Fin rows) : Result (Tape × ℕ)

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

Conv2D + pooling (ConvPool FFI)

def Runtime.Autograd.Cuda.Tape.conv2d {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (t : Tape) (kernelId biasId inputId : ℕ) : Result (Tape × ℕ)

Conv2D forward/backward via ConvPool FFI (single image, channels-first).

ConvTranspose2D (ConvPool FFI)

def Runtime.Autograd.Cuda.Tape.convTranspose2d {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (t : Tape) (kernelId biasId inputId : ℕ) : Result (Tape × ℕ)

ConvTranspose2D forward/backward via ConvPool FFI (single image, channels-first).

Generic naming wrappers

The CUDA tape exposes conv/max_pool/avg_pool/smooth_max_pool using the same names as the CPU tape. These dispatch to the ConvPool CUDA FFI entrypoints that take per-axis parameters as Array Nat (rank ≤ 8).

The *2d* wrappers remain as concise convenience names for the common rank-2 case.

def Runtime.Autograd.Cuda.Tape.conv {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} (t : Tape) (kernelId biasId inputId : ℕ) (hInC : inC ≠ 0) (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) : Result (Tape × ℕ)

N-D convolution (CUDA) via ConvPool FFI (rank ≤ 8).

def Runtime.Autograd.Cuda.Tape.convTranspose {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} (t : Tape) (kernelId biasId inputId : ℕ) (hInC : inC ≠ 0) (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) : Result (Tape × ℕ)

N-D transposed convolution (CUDA) via ConvPool FFI (rank ≤ 8).

def Runtime.Autograd.Cuda.Tape.maxPool2d {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

MaxPool2D via ConvPool FFI (no padding).

def Runtime.Autograd.Cuda.Tape.maxPool2dPad {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

MaxPool2D via ConvPool FFI (with symmetric padding).

def Runtime.Autograd.Cuda.Tape.maxPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

N-D max pooling (CUDA) via ConvPool FFI (rank ≤ 8).

def Runtime.Autograd.Cuda.Tape.smoothMaxPool2d {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape) (xId : ℕ) (beta : Float) : Result (Tape × ℕ)

Smooth max-pool2d (log-sum-exp surrogate) via ConvPool FFI (no padding).

def Runtime.Autograd.Cuda.Tape.smoothMaxPool2dPad {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape) (xId : ℕ) (beta : Float) : Result (Tape × ℕ)

Smooth max-pool2d (log-sum-exp surrogate) via ConvPool FFI (with symmetric padding).

def Runtime.Autograd.Cuda.Tape.smoothMaxPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (t : Tape) (xId : ℕ) (beta : Float) : Result (Tape × ℕ)

N-D smooth max pooling (CUDA) via ConvPool FFI (rank ≤ 8).

def Runtime.Autograd.Cuda.Tape.avgPool2d {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

AvgPool2D via ConvPool FFI (no padding).

def Runtime.Autograd.Cuda.Tape.avgPool2dPad {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

AvgPool2D via ConvPool FFI (with symmetric padding).

def Runtime.Autograd.Cuda.Tape.avgPool {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

N-D average pooling (CUDA) via ConvPool FFI (rank ≤ 8).

Normalization

def Runtime.Autograd.Cuda.Tape.layerNorm {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (t : Tape) (xId gammaId betaId : ℕ) : Result (Tape × ℕ)

LayerNorm over the last dimension for (seqLen, embedDim) buffers.

This implementation uses the standard stable formulas and is expressed in terms of existing CUDA kernels (axis reductions + broadcasts + pointwise ops).

def Runtime.Autograd.Cuda.Tape.batchnormChannelFirst {channels height width : ℕ} (h_c : channels > 0) (h_h : height > 0) (h_w : width > 0) (t : Tape) (xId gammaId betaId : ℕ) : Result (Tape × ℕ)

BatchNorm for a single channel-first image (C,H,W) (no batch axis).

We normalize per-channel across the spatial dimension H*W, reusing the same math as layer-norm by treating the buffer as a (channels, height*width) matrix.

Softmax (last axis, row folding)

We implement softmax along the last axis by folding all leading dimensions into one rows axis. This covers:

• 2D softmax ((rows, cols)),
• 3D batched softmax ((batch, rows, cols)) by folding batch*rows into rows.

def Runtime.Autograd.Cuda.Tape.softmax {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

def Runtime.Autograd.Cuda.Tape.logSoftmax {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Stable log-softmax along the last axis, implemented directly on CUDA buffers.

Multi-head self-attention

Forward structure matches Spec.MultiHeadAttention.forward:

1. Q = x @ Wq, K = x @ Wk, V = x @ Wv
2. reshape to heads (numHeads, n, headDim)
3. attention per head (batched): softmax(Q Kᵀ / sqrt(headDim)) @ V
4. combine heads, then output projection @ Wo

Masking:

• If mask is provided, we upload it as a float32 {0,1} matrix and apply the same semantics as the spec: masked logits are replaced by -1000.0 before softmax, and gradients through masked entries are zeroed.
• This incurs a host-to-device copy for the mask (since the mask is a host Tensor Bool).

def Runtime.Autograd.Cuda.Tape.multiHeadAttention {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (t : Tape) (wqId wkId wvId woId xId : ℕ) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) (useFlash : Bool := true) : Result (Tape × ℕ)