CUDA Tape Operations: Shared Helpers

def Runtime.Autograd.LibTorch.Tape.validateU32Dimensions (opName : String) (dims : Array ℕ) : Result Unit

Reject shape metadata containing a dimension outside the CUDA UInt32 ABI.

def Runtime.Autograd.LibTorch.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.LibTorch.Tape.broadcastScalarToShape (g : Buffer) (outShape : Spec.Shape) : Buffer

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

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

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

def Runtime.Autograd.LibTorch.Tape.rowSoftmaxForward (x : Buffer) (rows cols : UInt32) : Buffer.WithWorkspace

Row-wise stable softmax.

The returned WithWorkspace owns the buffers used to compute the stable formula. Backward needs only the output, so callers can release the workspace as soon as they have consumed the forward result.

def Runtime.Autograd.LibTorch.Tape.rowSoftmaxBwd (y dLdy : Buffer) (rows cols : UInt32) : Buffer

Row-wise softmax VJP, centering the upstream values at a maximal probability before taking their weighted mean. This retains small derivatives at saturated outputs.

def Runtime.Autograd.LibTorch.Tape.rowLogSoftmaxForward (x : Buffer) (rows cols : UInt32) : Buffer.WithWorkspace

Row-wise stable log-softmax.

One maximal exponential supplies the exact unit term. The other exponentials are summed separately before log1p, retaining small tails that would disappear in log(1 + tail). Backward needs only the output; callers can release the workspace after the forward result.

def Runtime.Autograd.LibTorch.Tape.rowLogSoftmaxBwd (y dLdy : Buffer) (rows cols : UInt32) : Buffer

Row-wise log-softmax VJP. Separate the pivot contribution and use expm1 there to avoid subtracting two rounded unit probabilities. TorchLean still owns the backward formula.