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
• Blanchard, Higham and Higham, "Accurate Computation of the Log-Sum-Exp and Softmax Functions": https://arxiv.org/abs/1909.03469
• Barron, Gnanapragasam and Tschannen, "A Fix for Softmax and Cross-Entropy": https://arxiv.org/abs/2610.31642
• 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)

inductive Activation.Kind : Type

Activation functions with a parameter-free pointwise interpretation.

This type is shared by model specifications and public model builders. Keeping the choice in the specification layer prevents configuration strings from silently selecting the wrong semantics.

• relu : Kind

  Rectified linear unit, max(0, x).

• gelu : Kind

  GELU in its tanh approximation (Math.geluSpec). This is PyTorch's nn.GELU(approximate='tanh'), not the default erf-based nn.GELU().

• silu : Kind

  SiLU/Swish, x * sigmoid(x).

• tanh : Kind

  Hyperbolic tangent.

• sigmoid : Kind

  Logistic sigmoid.

@[instance_reducible]

instance Activation.instReprKind : Repr Kind

def Activation.instReprKind.repr : Kind → ℕ → Std.Format

@[instance_reducible]

instance Activation.instDecidableEqKind : DecidableEq Kind

Scalar activations

def Activation.Math.reluSpec {α : Type} [Zero α] [Max α] [BEq α] (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).

At zero we return the scalar zero directly. This matters for dual numbers: their equality compares primals, so this branch clears every tangent at the kink, matching reluDerivSpec. Elementwise maximum has a different convention and splits a tie equally. Every other input retains the backend's max result, including a floating-point NaN, which does not compare equal to zero.

@[simp]

theorem Activation.Math.relu_zero_branch_eq_max {α : Type} [LinearOrder α] [Zero α] (x : α) : (if (x == 0) = true then 0 else max x 0) = max x 0

The explicit zero test preserves maximum on ordered scalars with ordinary equality.

@[simp]

theorem Activation.Math.reluSpec_eq_max {α : Type} [LinearOrder α] [Zero α] (x : α) : reluSpec x = max x 0

On an ordered scalar with ordinary equality, ReLU is the usual maximum with zero.

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, evaluated with a nonpositive exponential argument:

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

For positive inputs we use this expression directly. For zero and negative inputs we use the equivalent ratio $\exp(x)/(1+\exp(x))$. Both denominators lie between 1 and 2 over the reals. In floating-point arithmetic this avoids forming exp(-x) when x is a large negative number. The same branch also matters for dual numbers: differentiating a quotient with an infinite denominator can produce NaN, even when the sigmoid value has rounded to zero.

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 evaluate the equivalent exponential ratio directly, rather than subtracting a rounded sigmoid value from one. This retains representable derivatives on both tails.

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)$.

The runtime squares a scaled reciprocal cosh instead of subtracting rounded squares from one.

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 axis-parametric tensor operation Activation.softmaxSpec.

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

Derivative of logisticSpec, expressed in output form.

def Activation.Math.eluSpec {α : Type} [Context α] (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 = α.

The negative branch uses the scalar context's compensated exponential difference, so small inputs do not disappear when exp x rounds to one. Differentiation carriers supply its chain rule.

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} [RatCast α] : α

The GELU tanh coefficient, rounded from the exact rational 44715 / 1000000.

Casting the fraction once avoids overflowing its denominator in small floating-point formats.

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

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

PyTorch analogy: torch.nn.functional.gelu(x, approximate="tanh"). This is not PyTorch's default nn.GELU(), which uses the exact erf form. This activation selects the tanh approximation; the separate pointwise erf operation depends on the scalar context's supported operations.

We evaluate the gate as sigmoid (2 * inner), the same real function as (1 + tanh inner) / 2. This avoids cancelling the negative tail when tanh rounds to minus one; it is not a different approximation to the Gaussian CDF.

def Activation.Math.geluDerivSpec {α : Type} [TorchLean.Storage α] [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.

On the negative branch, split the exponential and multiply by x before its second small factor. The product can remain representable after exp x, and hence sigmoid itself, underflows.

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

Derivative of Swish / SiLU.

The negative branch groups x + denominator before multiplying the split exponential factors. Forming sigmoid and its slope separately could underflow both before multiplying by x.

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

log(1 + x) with compensation for rounding in the addition. When 1 + x rounds to one, return x rather than losing the small term. Equality, not order, selects this case so the same formula also works for complex scalars. This is not a correct-rounding guarantee for the backend's logarithm.

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

exp(x) - 1, retaining small arguments when the exponential rounds to one.

Near zero, the logarithmic correction accounts for rounding in exp(x). Outside the unit neighborhood we use direct subtraction, avoiding a logarithm's branch ambiguity for complex scalars. The zero-exponential branch handles underflow. As with log1pSpec, accuracy still depends on the scalar implementation of exp and log.

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)$.

Using the same stable sigmoid keeps the derivative finite in both tails. Evaluating this formula over dual scalars also propagates the softplus second derivative without a large exponential.

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

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 : α) (ε : α := Context.defaultEpsilon) : α

Derivative of safeLogSpec.

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

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 : α) (ε : α := Context.defaultEpsilon) : α

Derivative of smoothAbsSpec.

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

Tensor-level tanh (pointwise).

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

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

Tensor-level ReLU (pointwise).

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

Tensor-level sigmoid (pointwise).

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

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

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

Tensor-level sigmoid derivative (pointwise).

def Activation.sigmoidBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (x seed : TorchLean.Tensor α s) : TorchLean.Tensor α s

Apply sigmoid's pointwise pullback while retaining the incoming gradient in range repairs.

def Activation.sigmoidOutputDerivSpec {α : Type} [TorchLean.Storage α] [One α] [Sub α] [Mul α] {s : Spec.Shape} (sigmoidOutput : TorchLean.Tensor α s) : TorchLean.Tensor α s

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

This output-only formula cannot recover a small slope once the activation has rounded to one. Use sigmoidDerivSpec on the preactivation when that value is available.

def Activation.tanhOutputDerivSpec {α : Type} [TorchLean.Storage α] [One α] [Sub α] [Mul α] {s : Spec.Shape} (tanhOutput : TorchLean.Tensor α s) : TorchLean.Tensor α s

Derivative of tanh from its cached output. At rounded endpoints this may lose a representable slope; tanhDerivSpec uses the preactivation instead.

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

Tensor-level tanh derivative (pointwise).

def Activation.tanhBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (x seed : TorchLean.Tensor α s) : TorchLean.Tensor α s

Apply tanh's pointwise pullback without separately rounding away a recoverable local slope.

def Activation.Kind.applySpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : Kind → TorchLean.Tensor α s → TorchLean.Tensor α s

Apply a parameter-free pointwise activation to a tensor.

def Activation.Kind.derivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : Kind → TorchLean.Tensor α s → TorchLean.Tensor α s

Apply the derivative selected by a parameter-free pointwise activation.

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} [TorchLean.Storage α] [Context α] {n : ℕ} (t : TorchLean.Tensor α [n.succ]) : TorchLean.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} [TorchLean.Storage α] [Context α] {n : ℕ} (t : TorchLean.Tensor α [n.succ]) : TorchLean.Tensor α [n.succ]

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

def Activation.Internal.normalizeWeights {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (values : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Normalize finite weights without requiring their sum to fit in the output format.

The ordinary left fold is retained unless it loses range or swallows a nonzero term. In that case, the complete sum and quotients are replayed over finite rational coefficient trees before rounding each output. This also retains derivatives carried by nested scalars. Nonfinite inputs and carriers without coefficient conversion keep the ordinary formula; a genuine zero sum is not repaired. Information lost while constructing the weights cannot be recovered here.

def Activation.softmaxVecSpec {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (t : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

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 maximum bounds the exponentials by one. The shared normalizer also handles an overflowing or stagnating denominator when finite coefficient conversion is available; shifting alone does not protect that reduction in low precision.

def Activation.Internal.softmaxInnermostSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : TorchLean.Tensor α s → TorchLean.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.

@[simp]

theorem Activation.unstack_softmaxInnermostSpec_matrix {α : Type} [TorchLean.Storage α] [Context α] {rows columns : ℕ} (tensor : TorchLean.Tensor α [rows, columns]) (row : Fin rows) : (Internal.softmaxInnermostSpec tensor).unstack row = softmaxVecSpec (tensor.unstack row)

Last-axis softmax on a matrix acts independently on each row.

@[noinline]

def Activation.Internal.softmaxPullbackCoefficients? {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (weights : TorchLean.Tensor α [n]) (gradient : Fin n → Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))) (origin : Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))) : Option (Fin n → Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α)))

Form a centered softmax pullback from finite, unrounded gradient coefficients.

Attention can supply the original contraction instead of encoding an already overflowed tensor. The caller retains its ordinary result if the original inputs cannot be encoded.

@[noinline]

def Activation.Internal.replaySoftmaxPullback? {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (weights ordinary : TorchLean.Tensor α [n]) (gradient : Fin n → Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))) (origin : Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))) : Option (TorchLean.Tensor α [n])

Decode a centered softmax pullback only after combining all finite coefficients.

def Activation.Internal.softmaxPullback {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (weights gradient : TorchLean.Tensor α [n]) (anchor : α) : TorchLean.Tensor α [n]

Apply a weighted, centered softmax Jacobian to one row.

Finite input coefficients are replayed only when the ordinary result loses range. In particular, the difference of two finite incoming gradients may overflow even when every final component fits. Replaying the whole contraction avoids both that intermediate and spurious 0 * ∞ at blocked weights. Nonfinite inputs and unsupported carriers retain the ordinary formula.

def Activation.Internal.softmaxInnermostBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.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). $$

Before taking the weighted mean, subtract the upstream value at a maximal logit from every upstream coordinate. This leaves the real VJP unchanged, avoids materializing a Jacobian, and retains small derivatives when the dominant probability has rounded to one. Constant upstream vectors become exactly zero before either multiplication or reduction.

def Activation.softmaxSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (x : TorchLean.Tensor α s) : TorchLean.Tensor α s

Numerically stable softmax along any tensor dimension.

The selected dimension is moved to the innermost position, where a private kernel computes each one-dimensional slice, and is then restored. This definition covers outer and interior dimensions without imposing a memory-layout convention on the mathematical tensor.

def Activation.softmaxBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (x dY : TorchLean.Tensor α s) : TorchLean.Tensor α s

Backward/VJP for softmax along any tensor dimension.

def Activation.logSoftmaxVecSpec {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (t : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Log-softmax on a length-n vector.

One maximal entry supplies the unit term of the normalizer. We sum the other exponentials separately and use log1pSpec, so a small tail is not lost in an addition to one. Exactly one entry is omitted, including when several entries tie for the maximum.

def Activation.Internal.logSoftmaxInnermostSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s

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

def Activation.Internal.logSoftmaxInnermostJvpSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.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. The pivot's contribution is separated with expm1Spec, retaining its small complementary probability in forward mode as well.

def Activation.Internal.logSoftmaxVecBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (y dY : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n]

Vector log-softmax VJP, keeping the pivot's small complementary probability.

For the largest log-probability, 1 - exp(y) is evaluated as -expm1(y). Its own upstream coordinate is separated from the remaining sum before multiplying. Thus a one-hot loss does not subtract two values that have already rounded to one. The other coordinates use the usual linear-time VJP. This formula also retains the meaning of arbitrary, unnormalized upstream values.

def Activation.Internal.logSoftmaxInnermostBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.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}. $$

We take the already-computed output y rather than the logits x. The vector implementation rearranges the pivot coordinate to avoid cancellation when its probability is close to one.

def Activation.logSoftmaxSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (x : TorchLean.Tensor α s) : TorchLean.Tensor α s

Numerically stable log-softmax along any tensor dimension.

def Activation.logSoftmaxJvpSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (y dx : TorchLean.Tensor α s) : TorchLean.Tensor α s

Forward-mode derivative of log-softmax along any tensor dimension.

def Activation.logSoftmaxBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (y dY : TorchLean.Tensor α s) : TorchLean.Tensor α s

Backward/VJP for log-softmax along any tensor dimension.

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

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

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

Tensor-level derivative of leaky ReLU (pointwise).

def Activation.eluSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : TorchLean.Tensor α s) (alpha : α) : TorchLean.Tensor α s

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

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

Tensor-level derivative of ELU (pointwise).

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

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

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

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

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

Tensor-level Swish / SiLU (pointwise).

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

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

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

Tensor-level softplus (pointwise).

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

Tensor-level derivative of softplus (pointwise).

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

Tensor-level safeLogSpec (pointwise).

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

Tensor-level derivative of safeLogSpec (pointwise).

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

Tensor-level smoothAbsSpec (pointwise).

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

Tensor-level derivative of smoothAbsSpec (pointwise).