CUDA Tape Operations: Normalization and Row Softmax

def Runtime.Autograd.LibTorch.Tape.crossEntropy {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (t : Tape SomeBuffer) (logitsId targetId : ℕ) (mean : Bool := true) : Result (Tape SomeBuffer × ℕ)

Record weighted logits cross-entropy and its two native numerical pullbacks.

def Runtime.Autograd.LibTorch.Tape.distill {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (t : Tape SomeBuffer) (studentId teacherId : ℕ) : Result (Tape SomeBuffer × ℕ)

Record mean teacher-logit cross-entropy. The existing tape owns the student pullback; ATen evaluates its numerical formula without native graph recording.

Normalization

@[inline]

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

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

The tape records one normalization operation and keeps TorchLean's usual VJP. The native primitive evaluates the forward formula and that VJP with LibTorch tensor operations; intermediate tensors remain internal to the primitive.

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.

@[inline]

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

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)

All leading dimensions are folded into one row axis; normalization acts independently on the last axis at any rank. Empty tensors retain their shape. A scalar softmax is one, and a scalar log-softmax is zero; both have zero cotangent.

@[inline]

def Runtime.Autograd.LibTorch.Tape.Internal.rowOpLast {s : Spec.Shape} (t : Tape SomeBuffer) (opName : String) (xId : ℕ) (scalarForward : Buffer → Buffer) (forward : Buffer → UInt32 → UInt32 → Buffer.WithWorkspace) (backward : Buffer → Buffer → UInt32 → UInt32 → Buffer) : Result (Tape SomeBuffer × ℕ)

Record a row operation whose VJP uses the saved output.

The scalar forward function runs only after input validation. Empty tensors are copied; for non-scalar, nonempty tensors, row-dimension validation precedes input lookup.

@[inline]

def Runtime.Autograd.LibTorch.Tape.softmaxLast {s : Spec.Shape} (t : Tape SomeBuffer) (xId : ℕ) : Result (Tape SomeBuffer × ℕ)

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

@[inline]

def Runtime.Autograd.LibTorch.Tape.topKMask {s : Spec.Shape} (t : Tape SomeBuffer) (k xId : ℕ) : Result (Tape SomeBuffer × ℕ)

The top-k mask along the last axis (Spec.MoE.topKMaskSpec), as a node that receives no gradient. A scalar is a row of one entry, kept when 0 < k.

@[inline]

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

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