Activation Specifications

Scalar activation functions and their chosen derivatives live in Activation.Math. Tensor operations map those definitions pointwise, except for shape-dependent operations such as softmax and log-softmax. The definitions are polymorphic over the scalar Context, allowing the same layer specification to be interpreted over runtime floats, exact scalars, or verification domains.

The formulas and conventions follow these references:

• PyTorch activations: https://pytorch.org/docs/stable/nn.functional.html
• PyTorch torch.softmax: https://pytorch.org/docs/stable/generated/torch.softmax.html
• ReLU: Vinod Nair and Geoffrey Hinton, "Rectified Linear Units Improve Restricted Boltzmann Machines" (ICML 2010)
• ELU: Djork-Arne Clevert et al., "Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)" (ICLR 2016)
• GELU: Dan Hendrycks and Kevin Gimpel, "Gaussian Error Linear Units (GELUs)" (arXiv:1606.08415)
• Swish / SiLU: Prajit Ramachandran et al., "Searching for Activation Functions" (arXiv:1710.05941)

Scalar activations

def Activation.Math.reluSpec {α : Type} [Zero α] [Max α] (x : α) : α

ReLU: $\operatorname{ReLU}(x)=\max(x,0)$.

PyTorch analogy: torch.nn.functional.relu.

This is the simplest nonlinearity we use throughout TorchLean because it stays meaningful across many scalar backends (including ones that do not support exp/log).

def Activation.Math.reluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (x : α) : α

A standard subgradient choice for ReLU:

$\frac{d}{dx}\operatorname{ReLU}(x)=1$ if $x>0$, and $0$ otherwise.

PyTorch analogy: autograd picks a subgradient at $x=0$; our spec commits to a concrete one to make "the derivative" a pure function.

The DecidableRel (· > ·) constraint reflects that this definition branches on $x>0$.

def Activation.Math.sigmoidSpec {α : Type} [Context α] (x : α) : α

Logistic sigmoid:

$\operatorname{sigmoid}(x)=1/(1+\exp(-x))$.

PyTorch analogy: torch.nn.functional.sigmoid (or torch.sigmoid).

def Activation.Math.sigmoidDerivSpec {α : Type} [Context α] (x : α) : α

Derivative of sigmoid:

$\operatorname{sigmoid}'(x)=\sigma(x)(1-\sigma(x))$.

We write it this way (in terms of $\sigma(x)$) because that is the form used in most AD systems and it avoids re-expanding the exponential expression.

def Activation.Math.tanhSpec {α : Type} [Context α] (x : α) : α

Hyperbolic tangent: tanh(x). PyTorch analogy: torch.tanh.

def Activation.Math.tanhDerivSpec {α : Type} [Context α] (x : α) : α

Derivative of tanh:

$\tanh'(x)=1-\tanh^2(x)$.

def Activation.Math.leakyReluSpec {α : Type} [Zero α] [Mul α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (x αₗ : α) : α

Leaky ReLU:

$\operatorname{leaky\_relu}(x;\alpha)=x$ if $x>0$, else $\alpha x$.

PyTorch analogy: torch.nn.functional.leaky_relu with negative_slope = α.

def Activation.Math.leakyReluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (x αₗ : α) : α

Derivative of leaky ReLU:

$\frac{d}{dx}\operatorname{leaky\_relu}(x;\alpha)=1$ if $x>0$, else $\alpha$.

def Activation.Math.sinhSpec {α : Type} [Context α] (x : α) : α

Sinh: sinh(x).

def Activation.Math.sinhDerivSpec {α : Type} [Context α] (x : α) : α

Sinh derivative: cosh(x).

def Activation.Math.coshSpec {α : Type} [Context α] (x : α) : α

Cosh: cosh(x).

def Activation.Math.coshDerivSpec {α : Type} [Context α] (x : α) : α

Cosh derivative: sinh(x).

def Activation.Math.logisticSpec {α : Type} [Context α] (x : α) : α

Logistic form written as $\exp(x)/(\exp(x)+1)$.

This is mathematically the same sigmoid function as sigmoidSpec; we keep it as logisticSpec because several scalar approximation proofs reason about this exp(x) numerator form directly.

Important naming choice: this is not called scalar softmax. A one-entry softmax is always 1; the real softmax API in TorchLean is the tensor-level Activation.softmaxSpec below.

def Activation.Math.logisticDerivSpec {α : Type} [Context α] (x : α) : α

Derivative of logisticSpec, expressed in output form.

def Activation.Math.eluSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [MathFunctions α] [Sub α] [Mul α] (x alpha : α) : α

ELU (Exponential Linear Unit):

$\operatorname{ELU}(x;\alpha)=x$ if $x>0$, else $\alpha(\exp(x)-1)$.

PyTorch analogy: torch.nn.functional.elu with alpha = α.

def Activation.Math.eluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [MathFunctions α] [Mul α] (x alpha : α) : α

Derivative of ELU:

$\operatorname{ELU}'(x;\alpha)=1$ if $x>0$, else $\alpha\exp(x)$.

def Activation.Math.geluTanhCoeff {α : Type} [Context α] : α

The rational coefficient 44715 / 1000000 in the standard tanh approximation to GELU.

def Activation.Math.geluSpec {α : Type} [Context α] (x : α) : α

GELU (approximate): the common tanh-based approximation used in many Transformer codebases.

PyTorch analogy: torch.nn.functional.gelu(x, approximate="tanh").

def Activation.Math.geluDerivSpec {α : Type} [Context α] (x : α) : α

GELU derivative for the tanh-based approximation.

def Activation.Math.swishSpec {α : Type} [Context α] (x : α) : α

Swish / SiLU:

$\operatorname{swish}(x)=x\operatorname{sigmoid}(x)$.

PyTorch analogy: torch.nn.functional.silu.

def Activation.Math.swishDerivSpec {α : Type} [Context α] (x : α) : α

Derivative of Swish / SiLU.

Written in terms of sigmoid(x) for the same reason as sigmoidDerivSpec: this is the form used by AD systems and is convenient to reuse in proofs.

def Activation.Math.softplusSpec {α : Type} [Context α] (x : α) : α

Softplus, evaluated without a large positive exponential:

$\operatorname{softplus}(x)=\log(1+\exp(x))$.

The positive branch uses the equivalent expression $x+\log(1+\exp(-x))$; this keeps finite floating-point inputs finite when exp(x) itself would overflow. The operation remains the one-argument, unit-scale softplus used throughout TorchLean.

PyTorch analogy: torch.nn.functional.softplus.

def Activation.Math.softplusDerivSpec {α : Type} [Context α] (x : α) : α

Derivative of softplus:

$\operatorname{softplus}'(x)=\operatorname{sigmoid}(x)$.

def Activation.Math.safeLogSpec {α : Type} [Context α] (x : α) (ε : α := Numbers.epsilon) : α

A smooth log surrogate:

$\operatorname{safe\_log}(x;\varepsilon) =\log(\operatorname{softplus}(x)+\varepsilon)$.

We use this when we want something "log-like" without having to carry side conditions about the input being strictly positive.

def Activation.Math.safeLogDerivSpec {α : Type} [Context α] (x : α) (ε : α := Numbers.epsilon) : α

Derivative of safeLogSpec.

def Activation.Math.smoothAbsSpec {α : Type} [Context α] (x : α) (ε : α := Numbers.epsilon) : α

A smooth absolute value surrogate:

$\operatorname{smooth\_abs}(x;\varepsilon)=\sqrt{x^2+\varepsilon}$.

Useful when you want an abs-like shape but keep differentiability at 0.

def Activation.Math.smoothAbsDerivSpec {α : Type} [Context α] (x : α) (ε : α := Numbers.epsilon) : α

Derivative of smoothAbsSpec.

def Activation.tanhSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s

Tensor-level tanh (pointwise).

PyTorch analogy: torch.tanh(t) or torch.nn.functional.tanh(t) applied elementwise.

def Activation.reluSpec {α : Type} [Zero α] [Max α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level ReLU (pointwise).

def Activation.sigmoidSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level sigmoid (pointwise).

def Activation.reluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level ReLU derivative (pointwise), using the scalar subgradient choice in Activation.Math.reluDerivSpec.

def Activation.sigmoidDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level sigmoid derivative (pointwise).

def Activation.sigmoidOutputDerivSpec {α : Type} [Context α] {s : Spec.Shape} (sigmoidOutput : Spec.Tensor α s) : Spec.Tensor α s

Derivative of sigmoid when the sigmoid output has already been computed.

Recurrent layers save gate activations during the forward pass, so their backward specs should use this shared helper instead of re-defining s * (1 - s) locally.

def Activation.tanhDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level tanh derivative (pointwise).

Proper (last‑axis) softmax on tensors

These are the shape‑aware softmax definitions used in attention / classification layers. They recurse over outer dimensions and apply a numerically‑stable softmax to the last axis.

def Activation.maxVecSpec {α : Type} [Context α] {n : ℕ} (t : Spec.Tensor α (Spec.Shape.dim n.succ Spec.Shape.scalar)) : Spec.Tensor α Spec.Shape.scalar

Maximum entry of a nonempty vector, returned as a scalar tensor.

The fold is seeded by the first coordinate rather than by a numeric sentinel. Consequently the result is one of the input coordinates for every linearly ordered scalar type. Softmax and log-softmax share this definition so their range-reduction convention cannot drift apart.

def Activation.maxShiftedExpVecSpec {α : Type} [Context α] {n : ℕ} (t : Spec.Tensor α (Spec.Shape.dim n.succ Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim n.succ Spec.Shape.scalar)

Max-shifted exponentials shared by stable softmax and log-softmax.

def Activation.softmaxVecSpec {α : Type} [Context α] {n : ℕ} (t : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)

Softmax on a length-n vector.

This is the "real" softmax, not the scalar logistic helper in Activation.Math.logisticSpec.

Numerical stability:

We implement the standard stabilized form $\operatorname{softmax}(x)_i=\exp(x_i-m)/\sum_j\exp(x_j-m)$, where $m=\max_i x_i$. Subtracting the max avoids overflow in typical floating-point backends, and it is also a nice canonical form to reference in proofs.

def Activation.softmaxSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s

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

PyTorch analogy: torch.softmax(x, dim=-1).

For s = .scalar we return 1 (there is only one coordinate). For higher-rank tensors we keep the outer structure and apply softmaxVecSpec at the last axis.

def Activation.softmaxBackwardSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s

Backward/VJP for last-axis softmax.

If $y=\operatorname{softmax}(x)$ and we are given an upstream gradient $\partial L/\partial y$, then for each last-axis slice:

$$ \frac{\partial L}{\partial x} =y\odot\left( \frac{\partial L}{\partial y} -\left\langle\frac{\partial L}{\partial y},y\right\rangle \right). $$

This is the standard Jacobian-vector product for softmax, written in a way that avoids materializing the full n×n Jacobian.

def Activation.logSoftmaxVecSpec {α : Type} [Context α] {n : ℕ} (t : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)

Log-softmax on a length-n vector.

def Activation.logSoftmaxSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s

Log-softmax along the last axis (recurses over outer dimensions).

def Activation.logSoftmaxJvpSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s

Forward-mode JVP for last-axis log-softmax.

If $y=\operatorname{logsoftmax}(x)$, then each last-axis slice has directional derivative

$dy=dx-\operatorname{replicate}(\langle\exp(y),dx\rangle)$.

Unlike the VJP below, the subtracted scalar is replicated uniformly across the slice; the softmax probabilities occur only inside the dot product. Taking the already-computed output y also avoids recomputing the stable forward pass.

def Activation.logSoftmaxBackwardSpec {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s

Backward/VJP for last-axis log-softmax.

If $y=\operatorname{logsoftmax}(x)$, then $\operatorname{softmax}(x)=\exp(y)$ and the vector-Jacobian product is

$$ \frac{\partial L}{\partial x} =\frac{\partial L}{\partial y} -\operatorname{softmax}(x)\sum_i\frac{\partial L}{\partial y_i}. $$

This is the same formula used by PyTorch's stable log_softmax backward path. We take the already-computed output y rather than the logits x, so runtime backends can avoid recomputing the max-shifted forward pass during backprop.

def Activation.leakyReluSpec {α : Type} [Zero α] [Mul α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (t : Spec.Tensor α s) (αₗ : α) : Spec.Tensor α s

Tensor-level leaky ReLU (pointwise). PyTorch analogy: torch.nn.functional.leaky_relu.

def Activation.leakyReluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Spec.Shape} (t : Spec.Tensor α s) (αₗ : α) : Spec.Tensor α s

Tensor-level derivative of leaky ReLU (pointwise).

def Activation.eluSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [MathFunctions α] [Sub α] [Mul α] {s : Spec.Shape} (t : Spec.Tensor α s) (alpha : α) : Spec.Tensor α s

Tensor-level ELU (pointwise). PyTorch analogy: torch.nn.functional.elu.

def Activation.eluDerivSpec {α : Type} [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [MathFunctions α] [Mul α] {s : Spec.Shape} (t : Spec.Tensor α s) (alpha : α) : Spec.Tensor α s

Tensor-level derivative of ELU (pointwise).

def Activation.geluSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level GELU (approximate, pointwise). PyTorch analogy: gelu(..., approximate="tanh").

def Activation.geluDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level derivative of tanh-approx GELU (pointwise).

def Activation.swishSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level Swish / SiLU (pointwise).

def Activation.swishDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level derivative of Swish / SiLU (pointwise).

def Activation.softplusSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level softplus (pointwise).

def Activation.softplusDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) : Spec.Tensor α s

Tensor-level derivative of softplus (pointwise).

def Activation.safeLogSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) (ε : α := Numbers.epsilon) : Spec.Tensor α s

Tensor-level safeLogSpec (pointwise).

def Activation.safeLogDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) (ε : α := Numbers.epsilon) : Spec.Tensor α s

Tensor-level derivative of safeLogSpec (pointwise).

def Activation.smoothAbsSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) (ε : α := Numbers.epsilon) : Spec.Tensor α s

Tensor-level smoothAbsSpec (pointwise).

def Activation.smoothAbsDerivSpec {α : Type} [Context α] {s : Spec.Shape} (t : Spec.Tensor α s) (ε : α := Numbers.epsilon) : Spec.Tensor α s

Tensor-level derivative of smoothAbsSpec (pointwise).

def Activation.activationGradientSpec {α : Type} [Context α] {s : Spec.Shape} (activation_deriv : Spec.Tensor α s → Spec.Tensor α s) (input grad_output : Spec.Tensor α s) : Spec.Tensor α s

A generic pointwise activation VJP helper.

Given:

• f' (as a tensor-level derivative function),
• the forward input x,
• and an upstream gradient $\partial L/\partial f(x)$,

this returns $\partial L/\partial x$ by the chain rule:

$\frac{\partial L}{\partial x} =\frac{\partial L}{\partial f(x)}\odot f'(x)$.

This matches how most PyTorch elementwise ops behave in backward: multiply upstream gradients by the pointwise derivative mask/value.