Proof-Linked Session: Basic Graph Operations

Graph-node ops (implemented by reusing Compiled.GraphM)

def Runtime.Autograd.Torch.Internal.SessionIR.runGraphM {α : Type} {Γ : List Spec.Shape} {β : Type} (m : Compiled.GraphM.MWith α NatEnv Γ β) (ss : List Spec.Shape) (g : Proofs.Autograd.Algebra.GraphData α NatEnv Γ ss) : Result (β × (ss' : List Spec.Shape) × Proofs.Autograd.Algebra.GraphData α NatEnv Γ ss')

Run a Compiled.GraphM computation against the current (ss, g) pair.

Compiled.GraphM is the builder monad used by the proof-friendly compiled pipeline; reusing it here ensures this eager-style API records the same typed IR that the compiler expects.

def Runtime.Autograd.Torch.Internal.SessionIR.commitGraphM {α : Type} (s : SessionIR α) {β : Type} (k : {Γ ss : List Spec.Shape} → Proofs.Autograd.Algebra.TList α Γ → NatEnv → Proofs.Autograd.Algebra.GraphData α NatEnv Γ ss → Result (β × SessionIRState α)) : IO β

Atomically apply a graph-building update to the session snapshot.

This is the central adapter used by each op wrapper below: it reads s.st, runs a builder that returns an updated SessionIRState, stores it back into s.st, and returns the op result.

def Runtime.Autograd.Torch.Internal.SessionIR.const {α : Type} (s : SessionIR α) {sh : Spec.Shape} [Zero α] [DecidableEq Spec.Shape] (v : Spec.Tensor α sh) (name : Option String := none) : IO (TensorRef α sh)

Record a constant tensor.

Subtlety: if no op nodes have been created yet (ss = []), we record const as a leaf to match the eager session's leaf-collection behavior. Once op nodes exist, we emit an explicit constant node so users can introduce literal constants mid-graph. PyTorch comparison: like torch.tensor(...) (a leaf) vs inserting a literal constant into the graph; constants are treated as non-requires-grad.

def Runtime.Autograd.Torch.Internal.SessionIR.add {α : Type} (s : SessionIR α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise addition a + b.

PyTorch comparison: torch.add(a, b) / the + operator.

def Runtime.Autograd.Torch.Internal.SessionIR.sub {α : Type} (s : SessionIR α) [Sub α] [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise subtraction a - b.

PyTorch comparison: torch.sub(a, b) / the - operator.

def Runtime.Autograd.Torch.Internal.SessionIR.mul {α : Type} (s : SessionIR α) [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise multiplication a * b.

PyTorch comparison: torch.mul(a, b) / the * operator.

def Runtime.Autograd.Torch.Internal.SessionIR.scale {α : Type} (s : SessionIR α) [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (c : α) : IO (TensorRef α sh)

Record scaling by a scalar constant: x * c.

PyTorch comparison: like x * c (where c is a Python scalar).

def Runtime.Autograd.Torch.Internal.SessionIR.abs {α : Type} (s : SessionIR α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise absolute value.

PyTorch comparison: torch.abs(x).

def Runtime.Autograd.Torch.Internal.SessionIR.detach {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Stop-gradient boundary.

Forward semantics: identity. Backward semantics: no gradient flows to the input. PyTorch comparison: x.detach().

def Runtime.Autograd.Torch.Internal.SessionIR.sqrt {α : Type} (s : SessionIR α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise square root.

PyTorch comparison: torch.sqrt(x).

def Runtime.Autograd.Torch.Internal.SessionIR.clamp {α : Type} (s : SessionIR α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (x : TensorRef α sh) (minVal maxVal : α) : IO (TensorRef α sh)

Record elementwise clamp to the interval [minVal, maxVal].

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

def Runtime.Autograd.Torch.Internal.SessionIR.max {α : Type} (s : SessionIR α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise maximum of a and b.

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

def Runtime.Autograd.Torch.Internal.SessionIR.min {α : Type} (s : SessionIR α) [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {sh : Spec.Shape} (a b : TensorRef α sh) : IO (TensorRef α sh)

Record elementwise minimum of a and b.

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

def Runtime.Autograd.Torch.Internal.SessionIR.matmul {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {m n p : ℕ} (a : TensorRef α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : TensorRef α (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : IO (TensorRef α (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar)))

Record 2D matrix multiplication.

PyTorch comparison: torch.matmul(a, b) for 2D tensors.

def Runtime.Autograd.Torch.Internal.SessionIR.bmm {α : Type} (s : SessionIR α) [Add α] [Mul α] [Zero α] [DecidableEq Spec.Shape] {batch m n p : ℕ} (a : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)))) (b : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar)))) : IO (TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))))

Record batched matrix multiplication.

PyTorch comparison: torch.bmm(a, b) for 3D tensors of shape (batch, m, n) and (batch, n, p).

def Runtime.Autograd.Torch.Internal.SessionIR.concatVectors {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {n m : ℕ} (a : TensorRef α (Spec.Shape.dim n Spec.Shape.scalar)) (b : TensorRef α (Spec.Shape.dim m Spec.Shape.scalar)) : IO (TensorRef α (Spec.Shape.dim (n + m) Spec.Shape.scalar))

Concatenate two 1D vectors along dimension 0.

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

def Runtime.Autograd.Torch.Internal.SessionIR.concatDim0 {α : Type} (s : SessionIR α) [Context α] [DecidableEq Spec.Shape] {n m : ℕ} {sh : Spec.Shape} (a : TensorRef α (Spec.Shape.dim n sh)) (b : TensorRef α (Spec.Shape.dim m sh)) : IO (TensorRef α (Spec.Shape.dim (n + m) sh))

Concatenate two tensors along dimension 0.

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

def Runtime.Autograd.Torch.Internal.SessionIR.sliceRange0 {α : Type} (s : SessionIR α) [Zero α] [DecidableEq Spec.Shape] {n : ℕ} {sh : Spec.Shape} (x : TensorRef α (Spec.Shape.dim n sh)) (start len : ℕ) (h : len + start ≤ n) : IO (TensorRef α (Spec.Shape.dim len sh))

Slice a tensor along dimension 0.

This returns x[start : start+len]. The proof argument h enforces bounds. PyTorch comparison: x[start:start+len] for tensors with a leading dimension.