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} : 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} : (TorchLean.Tensor β s₁ → TorchLean.Tensor γ s₂) → 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} : TorchLean.Tensor α s → m (Ref m α s)

  Record a fixed tensor value.

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

  Add tensors elementwise.

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

  Subtract tensors elementwise.

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

  Multiply tensors elementwise.

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

  Multiply every element by a scalar.

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

  Take the elementwise absolute value.

• sqrt {s : Spec.Shape} : 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} : Ref m α s → α → α → 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} : Ref m α s → Ref m α s → m (Ref m α s)

  Take the elementwise maximum.

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

  Take the elementwise minimum.

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

  Repeat values along axes allowed by the broadcasting witness.

• reshape {s₁ s₂ : Spec.Shape} : 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 : ℕ) : 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] : 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] : Ref m α (batchA.concat [mDim, nDim]) → Ref m α (batchB.concat [nDim, pDim]) → m (Ref m α (batch.concat [mDim, pDim]))

  Multiply matrices after broadcasting their batch prefixes to a common shape.

• concatLeadingAxis {nDim mDim : ℕ} {s : Spec.Shape} : Ref m α (s.prependDim nDim) → Ref m α (s.prependDim mDim) → m (Ref m α (s.prependDim (nDim + mDim)))

  Concatenate tensors along their leading axis.

• sliceLeadingAxisRange {nDim : ℕ} {s : Spec.Shape} (start len : ℕ) (h : start + len ≤ nDim) : 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]} : 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]} : 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} : Ref m α s → m (Ref m α s)

  Replace negative elements with zero.

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

  Apply the logistic sigmoid elementwise.

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

  Apply the hyperbolic tangent elementwise.

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

  Apply tanh-approximate GELU as one backend primitive.

  The formula is 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x^3))). Keeping it primitive avoids building temporary tensors for each term. Backends must match Activation.geluSpec and Activation.geluDerivSpec.

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

  Apply softmax over the final tensor dimension.

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

  Apply stable log-softmax over the final tensor dimension.

  Use x - max(x) - log(sum(exp(x - max(x)))). Arbitrary-axis operations move their selected axis to the final position before calling this primitive.

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

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

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

  Take the elementwise exponential.

• sin {s : Spec.Shape} : 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} : Ref m α s → m (Ref m α s)

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

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

  Take the elementwise logarithm.

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

  Take the elementwise reciprocal.

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

  Keep the value and stop gradients through this reference.

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

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

  Positive epsilon keeps the logarithm's argument positive, including when a floating-point softplus rounds to zero. The derivative is sigmoid(x) / (softplus(x) + epsilon), with epsilon held fixed.

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

  Sum every element into a scalar.

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

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

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

  Apply weight * input + bias to one input vector.

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

  Average the squared elementwise difference between prediction and target.

• 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.

• multiHeadAttention {n numHeads dModel headDim : ℕ} (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 α [n, dModel] → Option (TorchLean.Tensor Bool [n, n]) → m (Ref m α [n, dModel])

  Apply multi-head self-attention to one sequence, using the optional mask.

• batchedMultiHeadAttention {batch n numHeads dModel headDim : ℕ} (hBatch : batch ≠ 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 α [batch, n, dModel] → Option (TorchLean.Tensor Bool [n, n]) → m (Ref m α [batch, n, dModel])

  Multi-head self-attention with an explicit leading batch axis.

  Its mathematical meaning is the leading-axis map of multiHeadAttention; implementations may execute the samples together, but may not change the mask convention or the per-sample forward/VJP semantics.

• conv {d inC outC : ℕ} {kernel stride padding inSpatial : TorchLean.Tensor ℕ [d]} : Ref m α (Spec.Shape.ofList (outC :: inC :: kernel.to (List ℕ))) → Ref m α [outC] → 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]} : Ref m α (Spec.Shape.ofList (inC :: outC :: kernel.to (List ℕ))) → Ref m α [outC] → 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} : 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.

• spectralConv1dRfftNative? : Option ({grid width modes : ℕ} → Ref m α [grid, width] → Ref m α [modes, width, width] → Ref m α [modes, width, width] → m (Option (Ref m α [grid, width])))

  Optional one-sided spectral convolution. Both weight tensors use [modes, input, output]; the transform is unnormalized and its real inverse divides by grid.

@[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.