GraphM Elementwise And Scalar Ops

Arithmetic, activations, scalar reductions, and MSE loss builders for proof-compiled graphs.

JVP vs VJP in this module

Each compiled node stores both:

• vjp: reverse-mode vector-Jacobian product (used by backprop), and
• jvp: forward-mode Jacobian-vector product (directional derivative).

The .compiled runtime path is primarily exercised via reverse-mode (VJP) and compilation to the eager tape. Basic elementwise/bilinear ops provide real JVP rules, shape-structural ops (for example slice/concat) apply the same transformation to the tangent, and heavier ops should expose named spec-layer JVP helpers before being wired here. Reverse-only ops it must be listed in reverseOnlyJvpOps and call unsupportedJvp rather than returning a silent zero tangent.

Forward-mode coverage is expanded by adding concrete jvp rules next to the corresponding forward and vjp definitions.

def Runtime.Autograd.Compiled.GraphM.add {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (a b : Var s) : MWith α Δ Γ (Var s)

Elementwise addition node (y = a + b).

PyTorch comparison: torch.add(a, b).

def Runtime.Autograd.Compiled.GraphM.sub {α Δ : Type} [Sub α] [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (a b : Var s) : MWith α Δ Γ (Var s)

Elementwise subtraction node (y = a - b).

PyTorch comparison: torch.sub(a, b).

def Runtime.Autograd.Compiled.GraphM.mul {α Δ : Type} [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (a b : Var s) : MWith α Δ Γ (Var s)

Elementwise multiplication node (y = a ⊙ b).

PyTorch comparison: torch.mul(a, b).

def Runtime.Autograd.Compiled.GraphM.square {α Δ : Type} [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Square x ↦ x ⊙ x.

def Runtime.Autograd.Compiled.GraphM.scale {α Δ : Type} [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) (c : α) : MWith α Δ Γ (Var s)

Scale a tensor by a scalar constant c (y = c * x).

PyTorch comparison: c * x / torch.mul(x, c).

def Runtime.Autograd.Compiled.GraphM.abs {α : Type} [Context α] [DecidableEq Spec.Shape] [DecidableRel fun (x1 x2 : α) => x1 > x2] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise absolute value.

PyTorch comparison: torch.abs(x).

def Runtime.Autograd.Compiled.GraphM.sqrt {α : Type} [Context α] [DecidableEq Spec.Shape] [DecidableRel fun (x1 x2 : α) => x1 > x2] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise square root.

PyTorch comparison: torch.sqrt(x).

def Runtime.Autograd.Compiled.GraphM.clamp {α : Type} [Context α] [DecidableEq Spec.Shape] [DecidableRel fun (x1 x2 : α) => x1 > x2] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) (minVal maxVal : α) : MWith α Δ Γ (Var s)

Elementwise clamp to [minVal, maxVal].

PyTorch comparison: torch.clamp(x, min=minVal, max=maxVal).

def Runtime.Autograd.Compiled.GraphM.max {α : Type} [Context α] [DecidableEq Spec.Shape] [DecidableRel fun (x1 x2 : α) => x1 > x2] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (a b : Var s) : MWith α Δ Γ (Var s)

Elementwise maximum.

At ties we split the gradient equally (0.5 / 0.5), matching the tie-handling documented in the eager tape (NN.Runtime.Autograd.Engine.Core).

PyTorch comparison: torch.maximum(a, b).

def Runtime.Autograd.Compiled.GraphM.min {α : Type} [Context α] [DecidableEq Spec.Shape] [DecidableRel fun (x1 x2 : α) => x1 > x2] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (a b : Var s) : MWith α Δ Γ (Var s)

Elementwise minimum.

At ties we split the gradient equally (0.5 / 0.5).

PyTorch comparison: torch.minimum(a, b).

def Runtime.Autograd.Compiled.GraphM.relu {α : Type} [Mul α] [Add α] [Zero α] [Max α] [One α] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise ReLU.

PyTorch comparison: torch.nn.functional.relu(x).

def Runtime.Autograd.Compiled.GraphM.sigmoid {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise sigmoid. PyTorch comparison: torch.sigmoid(x).

def Runtime.Autograd.Compiled.GraphM.tanh {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise tanh. PyTorch comparison: torch.tanh(x).

def Runtime.Autograd.Compiled.GraphM.gelu {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise tanh-approximate GELU.

This is one graph node with the specification-level derivative, rather than an expansion into temporary pointwise nodes. The smaller graph is important for large Transformer activations while retaining the same JVP and VJP meaning.

def Runtime.Autograd.Compiled.GraphM.softmax {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Softmax along the last axis (recursing over outer dimensions).

PyTorch comparison: torch.softmax(x, dim=-1).

def Runtime.Autograd.Compiled.GraphM.logSoftmax {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Stable log-softmax along the last axis.

This is a primitive in the compiled graph, not the composition log ∘ softmax, so proof/IR execution and eager CUDA share the same PyTorch-style numerical contract.

def Runtime.Autograd.Compiled.GraphM.softplus {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise softplus. PyTorch comparison: torch.nn.functional.softplus(x).

def Runtime.Autograd.Compiled.GraphM.exp {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise exponential. PyTorch comparison: torch.exp(x).

def Runtime.Autograd.Compiled.GraphM.log {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise natural logarithm. PyTorch comparison: torch.log(x).

def Runtime.Autograd.Compiled.GraphM.inv {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var s)

Elementwise reciprocal x ↦ 1/x. PyTorch comparison: torch.reciprocal(x).

def Runtime.Autograd.Compiled.GraphM.safeLog {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) (ε : α := Numbers.epsilon) : MWith α Δ Γ (Var s)

Elementwise numerically-stable log (uses an internal ε).

PyTorch comparison: commonly written torch.log(x + eps).

def Runtime.Autograd.Compiled.GraphM.sum {α : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (x : Var s) : MWith α Δ Γ (Var Spec.Shape.scalar)

Reduce-sum over all entries, producing a scalar.

PyTorch comparison: torch.sum(x).

def Runtime.Autograd.Compiled.GraphM.mseLoss {α : Type} [Add α] [Sub α] [Mul α] [Div α] [Zero α] [One α] [Coe ℕ α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (yhat target : Var s) : MWith α Δ Γ (Var Spec.Shape.scalar)

Mean-squared error loss with "mean" reduction, producing a scalar.

PyTorch comparison: torch.nn.functional.mse_loss(yhat, target, reduction=\"mean\").

def Runtime.Autograd.Compiled.GraphM.linear {α Δ : Type} [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {inDim outDim : ℕ} (w : Var (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b : Var (Spec.Shape.dim outDim Spec.Shape.scalar)) (x : Var (Spec.Shape.dim inDim Spec.Shape.scalar)) : MWith α Δ Γ (Var (Spec.Shape.dim outDim Spec.Shape.scalar))

Affine layer y = W x + b in the compiled graph.

PyTorch comparison: torch.nn.functional.linear / torch.nn.linear.

The JVP is the usual product rule: d(Wx+b) = dW*x + W*dx + db.

def Runtime.Autograd.Compiled.GraphM.matmul {α Δ : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {m n p : ℕ} (a : Var (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Var (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : MWith α Δ Γ (Var (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar)))

Matrix multiplication ((m×n) @ (n×p) → (m×p)).

PyTorch comparison: torch.matmul.

The JVP is the bilinear product rule d(A @ B) = dA @ B + A @ dB.

def Runtime.Autograd.Compiled.GraphM.bmm {α Δ : Type} [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {batch m n p : ℕ} (a : Var (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)))) (b : Var (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))))

Batched matrix multiplication (batch×m×n with batch×n×p).

PyTorch comparison: torch.bmm.

The JVP is the batched bilinear product rule d(A @ B) = dA @ B + A @ dB.

def Runtime.Autograd.Compiled.GraphM.concatVectors {α Δ : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n m : ℕ} (a : Var (Spec.Shape.dim n Spec.Shape.scalar)) (b : Var (Spec.Shape.dim m Spec.Shape.scalar)) : MWith α Δ Γ (Var (Spec.Shape.dim (n + m) Spec.Shape.scalar))

Concatenate two vectors (dim-0 concat).

PyTorch comparison: torch.cat([a, b], dim=0) for 1D tensors.

def Runtime.Autograd.Compiled.GraphM.concatLeadingAxis {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n m : ℕ} {s : Spec.Shape} (a : Var (Spec.Shape.dim n s)) (b : Var (Spec.Shape.dim m s)) : MWith α Δ Γ (Var (Spec.Shape.dim (n + m) s))

Concatenate along the leading dimension (dim=0) for tensors of shape .dim n s.

PyTorch comparison: torch.cat([a, b], dim=0).

def Runtime.Autograd.Compiled.GraphM.sliceLeadingAxisRange {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n : ℕ} {s : Spec.Shape} (x : Var (Spec.Shape.dim n s)) (start len : ℕ) (h : len + start ≤ n) : MWith α Δ Γ (Var (Spec.Shape.dim len s))

Slice a contiguous range along dim=0.

PyTorch comparison: x[start : start+len] for tensors where the leading dimension is indexed.