GraphM Shape And Indexing Ops

Reshape, transpose, broadcast, reduction, gather, and scatter builders for proof-compiled graphs.

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

Flatten a tensor to a 1D vector (preserving total size).

PyTorch comparison: torch.flatten(x) (for a single tensor value).

def Runtime.Autograd.Compiled.GraphM.reshape {α Δ : Type} [Inhabited α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (x : Var s₁) (h : s₁.size = s₂.size) : MWith α Δ Γ (Var s₂)

Reshape a tensor, given a proof that the total sizes match.

PyTorch comparison: torch.reshape(x, new_shape).

def Runtime.Autograd.Compiled.GraphM.transpose2d {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {m n : ℕ} (x : Var (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : MWith α Δ Γ (Var (Spec.Shape.dim n (Spec.Shape.dim m Spec.Shape.scalar)))

Transpose a 2D matrix. PyTorch comparison: x.transpose(0, 1) / x.T for matrices.

def Runtime.Autograd.Compiled.GraphM.transpose3dFirstToLast {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {a b c : ℕ} (x : Var (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim a Spec.Shape.scalar))))

Transpose a rank-3 tensor by moving the first axis to the last ((a,b,c) → (b,c,a)).

PyTorch comparison: x.permute(1, 2, 0).

def Runtime.Autograd.Compiled.GraphM.transpose3dLastToFirst {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {a b c : ℕ} (x : Var (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim c (Spec.Shape.dim a (Spec.Shape.dim b Spec.Shape.scalar))))

Transpose a rank-3 tensor by moving the last axis to the first ((a,b,c) → (c,a,b)).

PyTorch comparison: x.permute(2, 0, 1).

def Runtime.Autograd.Compiled.GraphM.transpose3dLastTwo {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {a b c : ℕ} (x : Var (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim a (Spec.Shape.dim c (Spec.Shape.dim b Spec.Shape.scalar))))

Swap the last two axes of a rank-3 tensor ((a,b,c) → (a,c,b)).

PyTorch comparison: x.transpose(1, 2) for a 3D tensor.

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

Swap two adjacent axes at a given nesting depth.

This is the compiled-graph analogue of the eager Tape.swapAdjacentAtDepth. PyTorch comparison: a permute that swaps two neighboring dimensions.

def Runtime.Autograd.Compiled.GraphM.broadcastTo {α Δ : Type} [Inhabited α] [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s₁ s₂ : Spec.Shape} (cb : s₁.CanBroadcastTo s₂) (x : Var s₁) : MWith α Δ Γ (Var s₂)

Broadcast x : s₁ to a larger shape s₂ (given a CanBroadcastTo witness).

PyTorch comparison: x.expand(...) / broadcasting semantics in elementwise ops.

def Runtime.Autograd.Compiled.GraphM.reduceSum {α Δ : Type} [Add α] [Zero α] [Inhabited α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis s] [wf : s.WellFormed] (x : Var s) : MWith α Δ Γ (Var (Spec.Tensor.shapeAfterSum s axis))

Reduce-sum along a given axis.

PyTorch comparison: torch.sum(x, dim=axis).

def Runtime.Autograd.Compiled.GraphM.reduceMean {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {s : Spec.Shape} (axis : ℕ) [valid : Spec.Shape.valid_axis_inst axis s] [wf : s.WellFormed] (x : Var s) : MWith α Δ Γ (Var (Spec.Tensor.shapeAfterSum s axis))

Reduce-mean along a given axis.

PyTorch comparison: torch.mean(x, dim=axis).

def Runtime.Autograd.Compiled.GraphM.gatherScalar {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n : ℕ} (x : Var (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) : MWith α Δ Γ (Var Spec.Shape.scalar)

Gather a single scalar from a vector at a known-in-bounds index.

PyTorch comparison: x[i] for a 1D tensor.

def Runtime.Autograd.Compiled.GraphM.gatherRow {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {rows cols : ℕ} (x : Var (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (i : Fin rows) : MWith α Δ Γ (Var (Spec.Shape.dim cols Spec.Shape.scalar))

Gather a row from a matrix at a known-in-bounds row index.

PyTorch comparison: x[i, :] for a 2D tensor.

def Runtime.Autograd.Compiled.GraphM.gatherScalarNat {α Δ : Type} [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n : ℕ} (x : Var (Spec.Shape.dim n Spec.Shape.scalar)) (i : ℕ) : MWith α Δ Γ (Var Spec.Shape.scalar)

Gather a scalar from a vector at a runtime Nat index.

If i is out of bounds we return 0 and propagate no gradient (matching the forward choice).

def Runtime.Autograd.Compiled.GraphM.gatherVecNat {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n k : ℕ} (x : Var (Spec.Shape.dim n Spec.Shape.scalar)) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : MWith α Δ Γ (Var (Spec.Shape.dim k Spec.Shape.scalar))

Gather a vector of length k from a length-n vector using an index tensor of Nats.

Out-of-bounds indices yield 0 at the corresponding output position.

PyTorch comparison: torch.gather for 1D inputs, with explicit bounds handling.

def Runtime.Autograd.Compiled.GraphM.gatherRowsNat {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {rows cols k : ℕ} (x : Var (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (idx : Spec.Tensor ℕ (Spec.Shape.dim k Spec.Shape.scalar)) : MWith α Δ Γ (Var (Spec.Shape.dim k (Spec.Shape.dim cols Spec.Shape.scalar)))

Gather k rows from a (rows×cols) matrix using an index vector of Nats.

Out-of-bounds indices yield a zero row.

PyTorch comparison: torch.index_select(x, dim=0, index=idx) with explicit bounds handling.

def Runtime.Autograd.Compiled.GraphM.scatterAddVec {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {n : ℕ} (x : Var (Spec.Shape.dim n Spec.Shape.scalar)) (v : Var Spec.Shape.scalar) (i : Fin n) : MWith α Δ Γ (Var (Spec.Shape.dim n Spec.Shape.scalar))

Scatter-add into a vector at a single in-bounds index.

scatter_add_vec x v i adds the scalar v into x[i].

PyTorch comparison: x.index_add_(dim=0, index=[i], source=[v]) (conceptually).

def Runtime.Autograd.Compiled.GraphM.scatterAddRow {α Δ : Type} [Add α] [Zero α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {rows cols : ℕ} (x : Var (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) (v : Var (Spec.Shape.dim cols Spec.Shape.scalar)) (i : Fin rows) : MWith α Δ Γ (Var (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar)))

Scatter-add into a matrix at a single in-bounds row index.

scatter_add_row x v i adds the row vector v into x[i, :].

PyTorch comparison: x.index_add_(dim=0, index=[i], source=v.unsqueeze(0)) (conceptually).