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)

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.

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

@[instance_reducible]

instance Activation.instReprKind : Repr Kind

@[instance_reducible]

instance Activation.instDecidableEqKind : DecidableEq Kind

@[reducible, inline]

abbrev Activation.Kind.geluTanh : Kind

Explicit spelling of the existing tanh GELU activation; .gelu remains compatible.

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

The sigmoid value uses the stable branch above, so negative tails retain their small positive derivatives until the exponential itself underflows. On the positive side, the derivative becomes zero once the sigmoid value rounds to one; this is the output-based derivative used by the graph VJP as well.

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 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} [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} [TorchLean.Storage α] [Context α] : α

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

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; TorchLean has no erf primitive, so only the tanh approximation is specified.

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.

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

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.sigmoidOutputDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (sigmoidOutput : TorchLean.Tensor α s) : TorchLean.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} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : TorchLean.Tensor α s) : TorchLean.Tensor α s

Tensor-level tanh derivative (pointwise).

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.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 max avoids overflow in typical floating-point backends, and it is also a nice canonical form to reference in proofs.

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.

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

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

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.

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.

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

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.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 α] [Zero α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [MathFunctions α] [Sub α] [Mul α] {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).

def Activation.vjpSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (activationDeriv : TorchLean.Tensor α s → TorchLean.Tensor α s) (input gradOutput : TorchLean.Tensor α s) : TorchLean.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.