CUDA Tape Operations: Shared Helpers

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

Reject shape metadata containing a dimension outside the CUDA UInt32 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 2 * sigmoid(2x) - 1.

Unlike (exp(2x)-1)/(exp(2x)+1), this form does not produce ∞/∞ and hence NaN for a large positive finite input. It also saturates to ±1 in both tails using the existing CUDA primitive-buffer operations.

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.rowSoftmaxForward (x : Buffer) (rows cols : UInt32) : Buffer.WithWorkspace

Row-wise stable softmax.

The returned WithWorkspace records the buffers used to compute the stable formula. The tape keeps those buffers only as long as the node may need them for backprop, then releases them explicitly.

def Runtime.Autograd.Cuda.Tape.rowHardMaskedSoftmaxForward (x mask : Buffer) (rows cols : UInt32) : Buffer.WithWorkspace

Row-wise hard-masked softmax.

mask is a {0,1} buffer with the same (rows, cols) shape as x; blocked entries contribute literal zero numerator. This matches Spec.hardMaskedSoftmaxSpec, not a finite additive sentinel.

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

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

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

Row-wise stable log-softmax.

This computes x - rowMax - log(sum(exp(x-rowMax))) directly, avoiding the less stable log(softmax(x)) route. As with softmax, the returned workspace buffers belong to the tape node until the backward pass has finished.

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

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