Autograd OpSpecs (spec layer)

This file defines small OpSpec building blocks (forward + VJP) for common tensor operations. The definitions are intentionally direct mathematical contracts and live purely in the spec layer.

The declarations follow a uniform pattern:

• Each operation below is an OpSpec: a pure forward plus a pure VJP backward.
• Most ops here package *Spec and derivative-spec definitions from NN/Spec/*.

Where this sits in TorchLean:

• NN.Spec.* files define pure denotational semantics: what tensors/layers mean.
• This file packages some of those pure definitions as unary OpSpecs: forward plus VJP.
• NN.Runtime.Autograd.* executes programs, tracks parameters, manages tapes/sessions, dispatches CUDA kernels, handles RNG, and lowers graphs for reusable typed execution.

This file adapts operations whose input-gradient VJP is naturally expressed as a single OpSpec. Larger multi-input or parameterized layers (convolution, attention, batchnorm, pooling, RNG) still have precise specs and runtime implementations, but their full backward state usually belongs in layer/runtime code rather than in this compact unary interface.

PyTorch analogy (approximately):

• A spec OpSpec is like a compact torch.autograd.Function where we write down the VJP directly.
• We do not model PyTorch's mutable ctx; the spec layer receives the input tensor x directly.

Elementwise lifting helpers

def Spec.liftElementwiseBackward {α : Type} [TorchLean.Storage α] [Mul α] {s : Shape} (df : α → α) : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.Tensor α s

Lift an elementwise backward using the chain rule: $\frac{\partial L}{\partial x}=f'(x)\frac{\partial L}{\partial y}$ pointwise.

This is the standard VJP pattern for elementwise ops.

PyTorch analogy: the "local backward" rule for a pointwise op multiplies by the derivative mask.

def Spec.reluOp {α : Type} [TorchLean.Storage α] [Mul α] [One α] [Zero α] [Max α] [BEq α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Shape} : OpSpec α s s

Elementwise ReLU OpSpec on any shape.

PyTorch analogy: torch.relu(x) / torch.nn.functional.relu(x).

def Spec.sigmoidOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise sigmoid OpSpec on any shape.

PyTorch analogy: torch.sigmoid(x).

def Spec.tanhOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise tanh OpSpec on any shape.

PyTorch analogy: torch.tanh(x).

def Spec.softplusOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise softplus OpSpec on any shape.

PyTorch analogy: torch.nn.functional.softplus(x).

def Spec.siluOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise SiLU (also called Swish) OpSpec on any shape.

PyTorch analogy: torch.nn.functional.silu(x).

def Spec.eluOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (eluAlpha : α) : OpSpec α s s

Elementwise ELU OpSpec on any shape.

def Spec.geluOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise tanh-approximate GELU OpSpec on any shape.

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

def Spec.sinhOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise hyperbolic sine OpSpec.

def Spec.coshOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise hyperbolic cosine OpSpec.

def Spec.softmaxOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] : OpSpec α s s

Softmax OpSpec along an explicitly selected tensor dimension.

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

def Spec.logSoftmaxOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] : OpSpec α s s

Stable log-softmax OpSpec along an explicitly selected tensor dimension.

Backward recomputes the forward output so the VJP uses the same axis-parametric semantics. Runtime engines may cache that output instead.

Linear layers

def Spec.linearOp {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] [One α] {inDim outDim : ℕ} (m : LinearSpec α inDim outDim) : OpSpec α [inDim] [outDim]

Linear layer as an OpSpec: $y=Wx+b$.

This OpSpec only returns the input gradient $\partial L/\partial x$. Parameter gradients for $W$ and $b$ are not part of OpSpec (those live at the graph/runtime level).

PyTorch analogy: torch.nn.functional.linear forward, with autograd producing gradients for x, W, and b.

def Spec.binaryElemOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) (f dfdx : α → α → α) : OpSpec α s s

Generic elementwise binary OpSpec with captured right-hand tensor and d/dx.

This is a "closure style" op: we treat the RHS tensor as a captured constant and only return the VJP with respect to the LHS input.

PyTorch analogy: in a tape/graph, rhs is typically another node; here we are writing the "lhs-only" derivative for convenience.

def Spec.scaleOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (c : α) : OpSpec α s s

Scale by constant scalar.

PyTorch analogy: x * c where c is a scalar constant.

Unary elementwise ops

def Spec.negOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Negation (-x).

def Spec.absOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Absolute value (uses signSpec for the subgradient).

PyTorch analogy: torch.abs(x). At $x=0$ we pick the subgradient $0$.

def Spec.smoothAbsOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (ε : α := Context.defaultEpsilon) : OpSpec α s s

Smooth absolute value (a differentiable surrogate for abs).

This is useful when you want to avoid a kink at 0 in optimization. PyTorch analogy: there is no single canonical smoothAbs, but it is similar in spirit to $\sqrt{x^2+\varepsilon}$-style smoothings.

def Spec.expOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise exp.

PyTorch analogy: torch.exp(x).

def Spec.logOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise natural logarithm.

Domain discipline: this is the raw mathematical/PyTorch-style rule. The VJP divides by x, so callers should use it only when the input is strictly positive. Runtime backends are allowed to reject nonpositive inputs rather than silently manufacture a gradient. Use safeLogOp when the intended model is the smooth surrogate $\log(\operatorname{softplus}(x)+\varepsilon)$.

PyTorch analogy: torch.log(x).

def Spec.safeLogOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (ε : α := Context.defaultEpsilon) : OpSpec α s s

Elementwise smooth logarithm surrogate, $\log(\operatorname{softplus}(x)+\varepsilon)$.

For $\varepsilon>0$ this is defined on every real input. Its VJP multiplies by $\operatorname{sigmoid}(x)/(\operatorname{softplus}(x)+\varepsilon)$.

PyTorch expression: torch.log(torch.nn.functional.softplus(x) + eps).

def Spec.sqrtOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise square root.

The forward pass evaluates sqrt (max x 0). The backward pass uses the local factor $1/(2\sqrt{x})$ when x > 0, and zero otherwise, then multiplies that factor by the upstream value. Checked scalar backends replay the complete weighted quotient if an intermediate loses range, using the root rounded by that backend. Over the reals this gives a zero VJP for $x\le0$. On IEEE backends, a zero factor multiplied by an infinite or NaN upstream value need not produce zero.

PyTorch analogy: torch.sqrt(x) on the positive region, with an explicit TorchLean subgradient choice outside the classical domain.

def Spec.squareOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise square, $x^2$.

def Spec.powOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise power with a captured RHS exponent tensor.

This is the VJP with respect to the base $x$ for $x^{\mathtt{rhs}}$. Domain restrictions are the usual ones for the scalar backend's power operation. The zero-exponent branch tests all differentiation coefficients, so a zero exponent primal does not erase mixed derivatives. For a constant exponent, checked range loss refines the complete weighted derivative instead of rounding an unrepresentable local factor before applying the incoming gradient.

def Spec.invOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise reciprocal, $1/x$.

Domain discipline: this is the raw reciprocal. Its VJP is $-1/x^2$, so callers should use it only when zero is excluded by the surrounding invariant. Use safeInvOp when the intended model is $1/(x+\varepsilon)$.

PyTorch analogy: torch.reciprocal(x) or 1 / x.

def Spec.safeInvOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α s s

Elementwise epsilon-shifted reciprocal, $1/(x+\varepsilon)$.

The forward pass delegates to safedivSpec with unit numerator, and the VJP uses the derivative of the same shifted expression. The shift alone does not exclude a zero denominator: over the reals, x = -Context.defaultEpsilon makes it zero.

PyTorch analogy: usually written manually as 1.0 / (x + eps).

Binary ops capturing a right-hand tensor

def Spec.addOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Add a captured RHS tensor, $x+\mathtt{rhs}$.

def Spec.subOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Subtract a captured RHS tensor, $x-\mathtt{rhs}$.

def Spec.mulOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise multiply by a captured RHS tensor.

def Spec.divOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise divide by a captured RHS tensor.

Domain discipline: this is the raw division rule. The VJP multiplies by 1/rhs, so callers should only use it when the captured denominator is known nonzero. Use safeDivOp when the intended model is x/(rhs+ε).

def Spec.safeDivOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise epsilon-shifted division by a captured RHS tensor, $x/(\mathtt{rhs}+\varepsilon)$. The shift does not guarantee a nonzero denominator or finite output.

PyTorch analogy: usually written manually as x / (rhs + eps).

def Spec.minOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise minimum with a captured right-hand tensor.

The backward pass multiplies the upstream value by a mask: one where the input is strictly smaller than rhs, zero where it is strictly larger, and half when neither comparison holds. The last case includes ties and unordered IEEE comparisons. This is the same mask as the two-input tape operation. Capturing rhs removes its gradient output; it does not transfer its half of a tied gradient to the remaining input. A zero mask can still produce NaN with a nonfinite upstream value.

Away from ties this is the usual derivative. At a tie, minimum is not differentiable, and the half-gradient is the convention used by the runtime. The half is a rational constant, so constructing it does not require a representable 2.

def Spec.maxOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (rhs : TorchLean.Tensor α s) : OpSpec α s s

Elementwise maximum with a captured right-hand tensor.

The backward pass multiplies the upstream value by a mask: one where the input is strictly larger than rhs, zero where it is strictly smaller, and half when neither comparison holds. As in minOp, the captured tensor keeps its share of the selected gradient even though this operation returns only the gradient with respect to the input.

The strict comparisons and their order match the two-input tape operation, including its fallback when neither comparison holds.

def Spec.leakyReluOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (αₗ : α) : OpSpec α s s

Leaky ReLU with slope parameter.

PyTorch analogy: torch.nn.functional.leaky_relu(x, negative_slope=alpha_l).

def Spec.clampOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (minVal maxVal : α) : OpSpec α s s

Clamp OpSpec with a fixed interval.

We choose the standard subgradient 1 strictly inside the interval and 0 at/outside the boundaries, matching clampDerivativeSpec.

Loss OpSpecs

def Spec.mseLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

MSE loss (returns a scalar), capturing the target.

def Spec.maeLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

MAE loss (returns a scalar), capturing the target.

def Spec.huberLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) (delta : α := 1) : OpSpec α s Shape.scalar

Huber loss (returns a scalar), capturing the target.

def Spec.crossEntropyLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : OpSpec α s Shape.scalar

Cross-entropy loss (returns a scalar), capturing the target distribution.

This is "cross-entropy between distributions": target is $p$, yhat is $q$. PyTorch analogy: closer to -(p * log(q)).mean() than to the logits-based torch.nn.functional.cross_entropy default.

def Spec.crossEntropyLogitsLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

Logits-based cross-entropy loss, capturing the target distribution.

def Spec.binaryCrossEntropyLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : OpSpec α s Shape.scalar

Binary cross-entropy loss on probability tensors, capturing the target tensor.

def Spec.cosineSimilarityLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) (epsilon : α := Context.defaultEpsilon) : OpSpec α s Shape.scalar

Cosine-similarity loss, capturing the target tensor.

def Spec.hingeLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

Hinge loss (returns a scalar), capturing the target.

def Spec.poissonLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

Poisson loss (returns a scalar), capturing the target.

def Spec.logCoshLossOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (target : TorchLean.Tensor α s) : OpSpec α s Shape.scalar

Log-cosh loss (returns a scalar), capturing the target.

Shape/structure ops

def Spec.reshapeOp {α : Type} [TorchLean.Storage α] {s t : Shape} (h : s.size = t.size) : OpSpec α s t

Reshape op (requires a size-equality proof).

PyTorch analogy: x.reshape(...) (or view), but here the shape relationship is explicit.

def Spec.swapAdjacentAxesOp {α : Type} [TorchLean.Storage α] {s : Shape} (depth : ℕ) : OpSpec α s (s.swapAdjacentAtDepth depth)

Swap adjacent axes at an arbitrary depth.

def Spec.constantOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (value : α) : OpSpec α s s

Fill a tensor with a constant (ignores input).

PyTorch analogy: torch.full_like(x, value) (but here we keep the input only to fit the OpSpec shape, and ignore its content).

def Spec.scalarToShapeOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : OpSpec α Shape.scalar s

Replicate a scalar to any shape; backward sums gradients back to a scalar.

PyTorch analogy: broadcasting a scalar in arithmetic, and in backward accumulating by sum.

def Spec.applyMaskOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (mask : TorchLean.Tensor Bool s) : OpSpec α s s

Apply boolean mask: keep where mask true, else set 0.

PyTorch analogy: torch.where(mask, x, 0).

def Spec.dropoutMaskedOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (p : α) (mask : TorchLean.Tensor Bool s) : OpSpec α s s

Masked inverted-dropout OpSpec with an explicit mask.

def Spec.matmulRightOp {α : Type} [TorchLean.Storage α] [Context α] {batchA batchB batch : Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (B : TorchLean.Tensor α (batchB.concat [n, p])) : OpSpec α (batchA.concat [m, n]) (batch.concat [m, p])

Matrix-rank multiplication with a captured right operand and broadcasted batch prefixes.

def Spec.matmulLeftOp {α : Type} [TorchLean.Storage α] [Context α] {batchA batchB batch : Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (A : TorchLean.Tensor α (batchA.concat [m, n])) : OpSpec α (batchB.concat [n, p]) (batch.concat [m, p])

Matrix-rank multiplication with a captured left operand and broadcasted batch prefixes.

def Spec.oneHotEmbeddingOp {α : Type} [TorchLean.Storage α] [Context α] {vocab embedDim seqLen : ℕ} (embedding : Embedding vocab embedDim α) : OpSpec α [seqLen, vocab] [seqLen, embedDim]

One-hot embedding as an OpSpec over the one-hot input. Parameter gradients stay outside OpSpec; this wrapper returns only dOneHot.

def Spec.sliceAxisRangeOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (start count : ℕ) (hRange : start + count ≤ s.axisSize axis) : OpSpec α s (s.replaceAxis axis count)

Slice a contiguous range along an arbitrary tensor axis.

Reductions and broadcasting

def Spec.reduceSumOp {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (axis : ℕ) [_valid : Shape.HasNonemptyAxis axis s] [_wf : s.WellFormed] : OpSpec α s (TorchLean.Tensor.shapeAfterSum s axis)

Reduce-sum along axis using a NonemptyAxis proof; backward broadcasts back.

PyTorch analogy: torch.sum(x, dim=axis) (with keepdim=false).

def Spec.binaryBroadcastOp {α : Type} [TorchLean.Storage α] [Context α] {s1 s2 t : Shape} (rhs : TorchLean.Tensor α s2) (cbx : s1.CanBroadcastTo t) (cby : s2.CanBroadcastTo t) (f dfdx : α → α → α) (reduceBack : TorchLean.Tensor α t → TorchLean.Tensor α s1) : OpSpec α s1 t

Generic broadcasting-aware binary OpSpec.

The caller supplies:

• explicit broadcast proofs (CanBroadcastTo) for both sides, and
• a reduceBack map that takes a gradient in the broadcasted shape t and reduces it back to the left shape s1.

PyTorch analogy: this is where PyTorch's implicit broadcasting rules and reduction-of-broadcasted gradients ("sum over broadcasted dimensions") happen. In TorchLean we keep those shape relations explicit.