CUDA Tape Operations: Normalization and Row Softmax

Normalization

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

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

The tape records one normalization operation and keeps TorchLean's usual VJP. A fused buffer primitive evaluates the forward formula and that VJP without materializing each reduction, broadcast, and pointwise intermediate as a separate device buffer.

epsilon is added to the variance before taking the square root. Backward reuses the normalized input and inverse standard deviation saved by forward, so both passes use the caller's value.

def Runtime.Autograd.Cuda.Tape.batchNorm {channels : ℕ} {spatial : Spec.Shape} (hWellFormed : (Spec.Shape.dim channels spatial).wellFormed) (t : Tape) (xId gammaId betaId : ℕ) (epsilon : Float := TorchLean.normalizationEpsilon) : Result (Tape × ℕ)

Batch normalization over every axis after the channel axis.

The spatial shape is folded to one contiguous dimension for the CUDA reduction. This is a view of the storage layout, not a rank-specific implementation.

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.softmaxLast {s : Spec.Shape} (t : Tape) (xId : ℕ) : Result (Tape × ℕ)

Record a last-axis softmax on the tape, returning the extended tape and the new node id.

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

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