Core Tape Activations and Losses

This file implements activation and loss tape nodes for the backend-independent autograd engine. Each node records the spec-layer forward value and a backward closure that computes the corresponding VJP contribution.

def Runtime.Autograd.Tape.crossEntropy {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (t : Tape (Spec.SomeTensor α)) (logitsId targetId : ℕ) (mean : Bool := true) : Result (Tape (Spec.SomeTensor α) × ℕ)

Record weighted logits cross-entropy and both pullbacks on the existing tape.

def Runtime.Autograd.Tape.distill {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (t : Tape (Spec.SomeTensor α)) (studentId teacherId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Cross-entropy from a fixed teacher's logits. The student uses the cancellation-resistant shared gradient; the teacher is a forward dependency, not a differentiable parent.

@[inline]

def Runtime.Autograd.Tape.sigmoid {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise logistic sigmoid activation.

This builds a tape node whose forward pass is Activation.sigmoidSpec, and whose backward pass applies Activation.sigmoidBackwardSpec. It uses the preactivation and retains the incoming gradient when repairing a local slope that would round to zero on its own.

PyTorch comparison: torch.sigmoid / torch.nn.functional.sigmoid. Reference: https://pytorch.org/docs/stable/generated/torch.sigmoid.html

@[inline]

def Runtime.Autograd.Tape.tanh {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise hyperbolic tangent activation.

Forward uses Activation.tanhSpec; backward uses Activation.tanhBackwardSpec, retaining the incoming gradient through the local slope's range repair.

PyTorch comparison: torch.tanh. Reference: https://pytorch.org/docs/stable/generated/torch.tanh.html

@[inline]

def Runtime.Autograd.Tape.silu {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise SiLU, using the shared split-exponential value and derivative specifications.

@[inline]

def Runtime.Autograd.Tape.gelu {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise tanh-approximate GELU.

The tape records GELU as one semantic operation. Its backward closure uses the derivative proved in NN.Proofs.Gradients.Activation; runtime backends may fuse the corresponding pointwise work without changing this tape-level rule.

@[inline]

def Runtime.Autograd.Tape.softmaxLast {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Softmax along the last axis (recursing over outer dimensions).

This is the tape primitive behind the general axis API after it moves the selected axis to the innermost position. Its backward pass avoids materializing an n×n Jacobian per slice.

PyTorch comparison: torch.softmax(x, dim=-1). Reference: https://pytorch.org/docs/stable/generated/torch.softmax.html

@[inline]

def Runtime.Autograd.Tape.logSoftmaxLast {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Stable log-softmax along the last axis.

The shared specification separates one maximal exponential from the sum and uses a compensated log1p for the remaining tail. Backward retains the same tail through expm1, rather than subtracting a rounded unit probability from one. This is the primitive used by cross-entropy.

@[inline]

def Runtime.Autograd.Tape.softplus {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise softplus activation.

Forward uses Activation.softplusSpec; backward uses Activation.softplusDerivSpec.

PyTorch comparison: torch.nn.functional.softplus. Reference: https://pytorch.org/docs/stable/generated/torch.nn.functional.softplus.html

@[inline]

def Runtime.Autograd.Tape.exp {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise exponential.

Forward uses expSpec; backward multiplies by exp(x) (pointwise), i.e. d/dx exp(x) = exp(x).

PyTorch comparison: torch.exp. Reference: https://pytorch.org/docs/stable/generated/torch.exp.html

@[inline]

def Runtime.Autograd.Tape.log {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise natural logarithm.

Forward uses logSpec; backward divides the incoming gradient by x (pointwise) (on its mathematical domain; this runtime does not model NaNs/Infs explicitly).

PyTorch comparison: torch.log. Reference: https://pytorch.org/docs/stable/generated/torch.log.html

@[inline]

def Runtime.Autograd.Tape.inv {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise reciprocal x ↦ 1/x.

Backward implements d/dx (x⁻¹) = -(x⁻¹)² (pointwise).

PyTorch comparison: torch.reciprocal. Reference: https://pytorch.org/docs/stable/generated/torch.reciprocal.html

@[inline]

def Runtime.Autograd.Tape.safeLog {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) (ε : α := Context.defaultEpsilon) : Result (Tape (Spec.SomeTensor α) × ℕ)

Elementwise log(softplus(x) + ε).

The forward value is Activation.safeLogSpec and the backward factor is Activation.safeLogDerivSpec, which is sigmoid(x) / (softplus(x) + ε). Over the reals with positive ε, the argument stays positive and the value approaches log ε as x tends to negative infinity. Executable arithmetic uses the chosen scalar representation; this operation does not validate ε or promise finite results for every input.

PyTorch comparison: torch.log(torch.nn.functional.softplus(x) + eps); PyTorch has no single safe_log primitive.

@[inline]

def Runtime.Autograd.Tape.sum {α : Type} [TorchLean.Storage α] [Add α] [Zero α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (xId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Reduce-sum over all entries, producing a scalar node.

Backward replicates the upstream scalar gradient to every entry of the input tensor (i.e. d/dx Σ_i x_i = 1 per coordinate).

PyTorch comparison: torch.sum(x) with dim=None. Reference: https://pytorch.org/docs/stable/generated/torch.sum.html

@[inline]

def Runtime.Autograd.Tape.mseLoss {α : Type} [TorchLean.Storage α] [Add α] [Sub α] [Mul α] [Div α] [Zero α] [One α] [NatCast α] {s : Spec.Shape} (t : Tape (Spec.SomeTensor α)) (yhatId targetId : ℕ) : Result (Tape (Spec.SomeTensor α) × ℕ)

Tape node for MSE loss with "mean" reduction.

The forward value is a scalar. The backward pass returns gradients for both inputs: dL/dyhat from Spec.mseDerivSpec, and dL/dtarget = -dL/dyhat.

PyTorch comparison: torch.nn.functional.mse_loss. Reference: https://pytorch.org/docs/stable/generated/torch.nn.functional.mse_loss.html