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.sigmoid {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

Elementwise logistic sigmoid activation.

This builds a tape node whose forward pass is Activation.sigmoidSpec, and whose backward pass multiplies the upstream gradient by Activation.sigmoidDerivSpec (i.e. σ(x) * (1 - σ(x)), pointwise).

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

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

Elementwise hyperbolic tangent activation.

Forward uses Activation.tanhSpec; backward uses Activation.tanhDerivSpec (pointwise derivative, usually 1 - tanh(x)^2).

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

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

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.

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

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

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

Stable log-softmax along the last axis.

Unlike log (softmax x), this uses the max-shifted x - max(x) - log(sum(exp(x - max(x)))) formulation. That matches the numerical contract of torch.nn.functional.log_softmax and is the right primitive for cross-entropy on logits.

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

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

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

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

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

Elementwise natural logarithm.

Forward uses logSpec; backward multiplies by 1/x (pointwise), i.e. d/dx log(x) = 1/x (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

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

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

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

Elementwise "safe log" that protects against log(0) by adding a small ε internally.

This uses Activation.safeLogSpec and Activation.safeLogDerivSpec. The exact behavior is controlled by the spec-layer definition; conceptually it is similar to log(x + ε) used in numerically-stable losses.

PyTorch comparison: commonly written as torch.log(x + eps) in user code (there is no single dedicated torch.safe_log primitive).

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

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

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

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