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.liftElementwise {α : Type} [TorchLean.Storage α] {s : Shape} (f : α → α) : TorchLean.Tensor α s → TorchLean.Tensor α s

Lift a scalar function to a tensor by pointwise map.

PyTorch analogy: most torch.* pointwise ops are vectorized elementwise maps.

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 multiplies by 1/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.

Domain discipline: TorchLean's spec-level sqrtSpec is total by clamping the forward value on nonpositive inputs. The VJP follows that convention and returns zero where $x\le0$, rather than introducing an artificial $1/\varepsilon$ spike.

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.

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

This is the safe API counterpart to invOp: the forward pass delegates to safedivSpec with unit numerator, and the VJP is the derivative of the same shifted expression.

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 safe division by a captured RHS tensor, $x/(\mathtt{rhs}+\varepsilon)$.

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 gives the input the full upstream gradient where it is strictly smaller than rhs, zero where it is strictly larger, and half at a tie. This is the same selected gradient 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.

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.

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 gives the input the full upstream gradient where it is strictly larger than rhs, zero where it is strictly smaller, and half at a tie. 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.dropoutInferenceOp {α : Type} [TorchLean.Storage α] {s : Shape} (p : α) : OpSpec α s s

Evaluation-mode dropout, which is the identity in both the forward and backward maps.

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.

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

Convenience: broadcasting-aware add with caller-provided reduction.

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

Convenience: broadcasting-aware mul with caller-provided reduction.