Backend Operation Interface

Ops supplies the tensor primitives for models shared by eager and typed graph execution. Ref and DataRef select the reference types of the current backend.

class Runtime.Autograd.Torch.Ops (m : Type → Type) (α : Type) [TorchLean.Storage α] [Context α] : Type 1

Tensor operations shared by eager and typed graph execution.

A model is polymorphic over this class. Eager instances execute and record a tape; typed graph instances build reusable graph data. Each instance must preserve the corresponding operator's forward and VJP semantics.

• Ref : Spec.Shape → Type

  A differentiable tensor handle, indexed by its shape.

• DataRef (β : Type) [TorchLean.Storage β] : Spec.Shape → Type

  Backend representation of non-differentiable tensor data.

• dataConst {β : Type} [TorchLean.Storage β] {s : Spec.Shape} (x : TorchLean.Tensor β s) : DataRef m α β s

  Lift fixed non-differentiable data into the backend representation.

• mapData {β γ : Type} [TorchLean.Storage β] [TorchLean.Storage γ] {s₁ s₂ : Spec.Shape} (f : TorchLean.Tensor β s₁ → TorchLean.Tensor γ s₂) (x : DataRef m α β s₁) : DataRef m α γ s₂

  Apply a pure transformation to non-differentiable data.

• updateBuffers? : Option ({ss : List Spec.Shape} → {s : Spec.Shape} → RefList (Ref m α) ss → Ref m α s → (TorchLean.TensorPack α ss → TorchLean.Tensor α s → IO (TorchLean.TensorPack α ss)) → m Unit)

  Observe a stateful layer's actual forward input and update its persistent buffers.

  Pure interpreters may omit this hook. Runtime interpreters retain the recorded state values for differentiation and write updated values only to their persistent, non-trainable storage.

• const {s : Spec.Shape} (t : TorchLean.Tensor α s) : m (Ref m α s)

  Record a fixed tensor value.

• observe : Ref m α Spec.Shape.scalar → m α

  Inspect a scalar for eager control flow. The comparison is not differentiated; the tape records only the selected branch. Static graph builders reject this dynamic observation.

• add {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Add tensors elementwise.

• sub {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Subtract tensors elementwise.

• mul {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Multiply tensors elementwise.

• div {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Divide tensors elementwise, recording the quotient-rule VJP at nonzero denominators.

• neg {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Negate each element directly, preserving the scalar's signed-zero convention.

• scale {s : Spec.Shape} (x : Ref m α s) (c : α) : m (Ref m α s)

  Multiply every element by a scalar.

• abs {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise absolute value.

• sqrt {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Compute sqrt(max(x, 0)) elementwise. The selected JVP and VJP are zero for x ≤ 0; at positive inputs they use 1 / (2 * sqrt(x)).

• clamp {s : Spec.Shape} (x : Ref m α s) (minVal maxVal : α) : m (Ref m α s)

  Clamp every element between the lower and upper bounds. Input gradients pass through the open interval and are zero at both endpoints and outside it. The bounds are fixed scalar settings.

• max {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Take the elementwise maximum.

• min {s : Spec.Shape} (a b : Ref m α s) : m (Ref m α s)

  Take the elementwise minimum.

• broadcastTo {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) (x : Ref m α s₁) : m (Ref m α s₂)

  Repeat values along axes allowed by the broadcasting witness.

• reshape {s₁ s₂ : Spec.Shape} (x : Ref m α s₁) (h : s₁.size = s₂.size) : m (Ref m α s₂)

  Change the shape while preserving scalar count and row-major order.

• swapAdjacentAtDepth {s : Spec.Shape} (depth : ℕ) (x : Ref m α s) : m (Ref m α (s.swapAdjacentAtDepth depth))

  Swap the two adjacent axes starting at depth.

• reduceSum {s : Spec.Shape} (axis : ℕ) [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] (x : Ref m α s) : m (Ref m α (TorchLean.Tensor.shapeAfterSum s axis))

  Sum along axis and remove it from the result shape.

• reduceMean {s : Spec.Shape} (axis : ℕ) [Spec.Shape.HasNonemptyAxis axis s] [s.WellFormed] : Ref m α s → m (Ref m α (TorchLean.Tensor.shapeAfterSum s axis))

  Average along axis and remove it from the result shape.

• select {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] : Ref m α s → Fin (s.axisSize axis) → m (Ref m α (s.eraseAxis axis))

  Select one position along axis, removing that axis from the result.

• indexSelect {s : Spec.Shape} (axis count : ℕ) [Spec.Shape.AxisInBounds axis s] : Ref m α s → DataRef m α (Fin (s.axisSize axis)) [count] → m (Ref m α (s.replaceAxis axis count))

  Gather positions along axis in the order given by the indices.

• scatterAdd {s : Spec.Shape} (axis count : ℕ) [Spec.Shape.AxisInBounds axis s] : Ref m α s → Ref m α (s.replaceAxis axis count) → DataRef m α (Fin (s.axisSize axis)) [count] → m (Ref m α s)

  Add source slices into the base at the selected positions along axis.

• matmul {batchA batchB batch : Spec.Shape} {mDim nDim pDim : ℕ} [batchA.BroadcastTo batch] [batchB.BroadcastTo batch] (a : Ref m α (batchA.concat [mDim, nDim])) (b : Ref m α (batchB.concat [nDim, pDim])) (scale : Option α := none) : m (Ref m α (batch.concat [mDim, pDim]))

  Multiply matrices after broadcasting their batch prefixes to a common shape. A fixed scale stays inside the contraction, including its local derivatives.

• softmaxMatmul {batch : Spec.Shape} {rows keys width : ℕ} : Ref m α (batch.concat [rows, keys]) → Ref m α (batch.concat [keys, width]) → Ref m α (batch.concat [rows, keys]) → (mask : optParam (Option (TorchLean.Tensor Bool [rows, keys])) none) → m (Ref m α (batch.concat [rows, width]))

  Normalize score rows, apply a differentiable multiplier, and mix value rows. The pullback retains the value contraction through normalization instead of requiring its intermediate cotangent to fit the scalar format. The optional mask excludes keys exactly.

• concat {nDim mDim : ℕ} {s : Spec.Shape} (a : Ref m α (s.prependDim nDim)) (b : Ref m α (s.prependDim mDim)) : m (Ref m α (s.prependDim (nDim + mDim)))

  Concatenate tensors along their leading axis.

• slice {nDim : ℕ} {s : Spec.Shape} (start len : ℕ) (h : start + len ≤ nDim) (x : Ref m α (s.prependDim nDim)) : m (Ref m α (s.prependDim len))

  Take len entries of the leading axis, starting at start.

• maxPool {d C : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} (x : Ref m α (Spec.Shape.ofList (C :: inSpatial.to (List ℕ)))) : m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ))))

  Apply spatial max pooling to one channels-first sample.

• avgPool {d C : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} (x : Ref m α (Spec.Shape.ofList (C :: inSpatial.to (List ℕ)))) : m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ))))

  Apply spatial average pooling to one channels-first sample, counting padding as zeros.

• smoothMaxPool {d C : ℕ} {inSpatial kernel stride padding : TorchLean.Tensor ℕ [d]} [DecidableEq α] : Ref m α (Spec.Shape.ofList (C :: inSpatial.to (List ℕ))) → α → m (Ref m α (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).to (List ℕ))))

  Apply a smooth approximation to spatial max pooling, with sharpness beta.

• relu {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Replace negative elements with zero.

• sigmoid {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply the logistic sigmoid elementwise.

• silu {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply SiLU with its input-based value and VJP, retaining representable negative tails.

• tanh {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply the hyperbolic tangent elementwise.

• gelu {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply tanh-approximate GELU as one recorded operation.

  The formula is 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3))). A backend may evaluate it through several numerical primitives and temporary tensors without recording separate tape nodes for the terms. The value and VJP follow Activation.geluSpec and Activation.geluDerivSpec.

• softmaxLast {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply softmax over the final tensor dimension.

• logSoftmaxLast {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply stable log-softmax over the final tensor dimension.

  Separate one maximal exponential and normalize the remaining tail with log1p. Arbitrary-axis operations move their selected axis to the final position before calling this primitive. Backward must preserve small complementary probabilities near one.

• softplus {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Apply log(1 + exp(x)) elementwise.

• exp {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise exponential.

• sin {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise sine of angles in radians. JVPs and VJPs multiply by cos(x), evaluated at the original input, so differentiation preserves the derivative's sign across periods.

• cos {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise cosine of angles in radians, with derivative -sin(x).

• log {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise logarithm.

• inv {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Take the elementwise reciprocal.

• detach {s : Spec.Shape} (x : Ref m α s) : m (Ref m α s)

  Keep the value and stop gradients through this reference.

• topKMask {s : Spec.Shape} (k : ℕ) (x : Ref m α s) : m (Ref m α s)

  One for the k largest entries of every row along the last axis and zero for the others, with ties going to the lower index (Spec.MoE.topKMaskSpec); a row with NaN can keep a different number of entries. The mask is constant away from ties and no gradient flows through it; a router learns through the gates the mask selects.

• safeLog {s : Spec.Shape} : Ref m α s → α → m (Ref m α s)

  Apply log(softplus(x) + epsilon) elementwise.

  The derivative is sigmoid(x) / (softplus(x) + epsilon), with epsilon held fixed. Over the reals, positive epsilon keeps the argument positive. Executable arithmetic needs epsilon to remain positive in its selected representation; this does not make nonfinite inputs finite.

• sum {s : Spec.Shape} (x : Ref m α s) : m (Ref m α Spec.Shape.scalar)

  Sum every element into a scalar.

• flatten {s : Spec.Shape} (x : Ref m α s) : m (Ref m α [s.size])

  Flatten all axes into one, preserving row-major order.

• linear {inDim outDim : ℕ} (w : Ref m α [outDim, inDim]) (b : Ref m α [outDim]) (x : Ref m α [inDim]) : m (Ref m α [outDim])

  Apply weight * input + bias to one input vector.

• mseLoss {s : Spec.Shape} (yhat target : Ref m α s) : m (Ref m α Spec.Shape.scalar)

  Average the squared elementwise difference between prediction and target.

• crossEntropy {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (logits target : Ref m α s) (mean : Bool := true) : m (Ref m α Spec.Shape.scalar)

  Cross-entropy from logits and target weights along axis. Both inputs receive gradients. The default mean divides by class slices, not target mass; mean := false keeps the sum.

• distill {s : Spec.Shape} (axis : ℕ) [Spec.Shape.AxisInBounds axis s] (student teacher : Ref m α s) : m (Ref m α Spec.Shape.scalar)

  Mean cross-entropy from teacher logits along axis. Only the student receives gradients; callers detach the teacher to clear forward-mode metadata as well.

• layerNorm {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : Ref m α [seqLen, embedDim]) (gamma beta : Ref m α [embedDim]) (epsilon : α := TorchLean.normalizationEpsilon) : m (Ref m α [seqLen, embedDim])

  Normalize each row over its embedding axis, then apply scale and bias.

  The denominator is sqrt(variance + epsilon). Forward and backward use the same epsilon; gradients are taken with respect to the input, scale, and bias while epsilon stays fixed.

• batchNorm {channels : ℕ} {sSpatial : Spec.Shape} (hWellFormed : (sSpatial.prependDim channels).wellFormed) (x : Ref m α (sSpatial.prependDim channels)) (gamma beta : Ref m α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) : m (Ref m α (sSpatial.prependDim channels))

  Normalize each channel over its spatial axes, then apply scale and bias.

  The denominator is sqrt(variance + epsilon). Forward and backward use the same epsilon; gradients are taken with respect to the input, scale, and bias while epsilon stays fixed.

• attention {n numHeads dModel headDim : ℕ} {batch : Option ℕ} (hBatch : batch.getD 1 ≠ 0) (h1 : n ≠ 0) : Ref m α [dModel, numHeads * headDim] → Ref m α [dModel, numHeads * headDim] → Ref m α [dModel, numHeads * headDim] → Ref m α [numHeads * headDim, dModel] → Ref m α (match batch with | none => [n, dModel] | some b => [b, n, dModel]) → Option (TorchLean.Tensor Bool [n, n]) → m (Ref m α (match batch with | none => [n, dModel] | some b => [b, n, dModel]))

  Self-attention, optionally over a leading batch dimension.

  batch := none selects one sequence; some b selects b independent sequences with shared projections and mask. Batching preserves the per-sequence forward and VJP semantics.

• conv {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (weight : Ref m α (Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ)))) (bias : Ref m α [outC]) (input : Ref m α (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ)))) : m (Ref m α (Spec.Shape.ofList (outC :: (Spec.convOutSpatial inSpatial kernel stride padding).to (List ℕ))))

  Apply spatial convolution to one channels-first sample.

• convTranspose {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} (weight : Ref m α (Spec.Shape.ofList (inC :: outC :: kernel.to (List ℕ)))) (bias : Ref m α [outC]) (input : Ref m α (Spec.Shape.ofList (inC :: inSpatial.to (List ℕ)))) : m (Ref m α (Spec.Shape.ofList (outC :: (Spec.convTransposeOutSpatial inSpatial kernel stride padding).to (List ℕ))))

  Apply spatial transpose convolution to one channels-first sample.

• randUniform {s : Spec.Shape} (seed : ℕ) : m (Ref m α s)

  Draw a uniform tensor from the seed and the backend's node or call index.

  No IO randomness is used, so a fixed graph can replay the draw.

• bernoulliMask {s : Spec.Shape} (keepProb : Ref m α Spec.Shape.scalar) (seed : ℕ) : m (Ref m α s)

  Draw a seeded mask with the scalar reference as its keep probability.

• rfft1dNative? : Option ({batch n : ℕ} → Ref m α [batch, n] → m (Option (Ref m α [batch, n / 2 + 1, 2])))

  Optional native packed real transform. Returning none leaves the operation to the generic differentiable implementation; a backend must not record any nodes before declining the call. The output stores real and imaginary components in its final axis.

• irfft1dNative? : Option ({batch n : ℕ} → Ref m α [batch, n / 2 + 1, 2] → m (Option (Ref m α [batch, n])))

  Optional normalized inverse of the packed real transform, with explicit output length.

• selectiveScanDiagNative? : Option ({seqLen state : ℕ} → Ref m α [state] → Ref m α [state] → Ref m α [seqLen, state] → Ref m α [state] → m (Option (Ref m α [seqLen, state])))

  Optional diagonal recurrence with coefficients shared across time and all four VJPs.

• selectiveScanDiagVarNative? : Option ({seqLen state : ℕ} → Ref m α [seqLen, state] → Ref m α [seqLen, state] → Ref m α [seqLen, state] → Ref m α [state] → m (Option (Ref m α [seqLen, state])))

  Optional diagonal recurrence with token-dependent coefficients and all four VJPs.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Ref {m : Type → Type} {α : Type} [TorchLean.Storage α] [Context α] [Ops m α] (s : Spec.Shape) : Type

Differentiable reference type of the current backend.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.DataRef {m : Type → Type} {α : Type} [TorchLean.Storage α] [Context α] [Ops m α] (β : Type) [TorchLean.Storage β] (s : Spec.Shape) : Type

Backend representation of a non-differentiable tensor.