GraphM Core

Typed variables, builder state, input binding, constants, random nodes, and detach for the proof-compiled graph authoring API.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.PGraphData (α Δ : Type) (Γ ss : List Spec.Shape) : Type

Shorthand for the underlying executable SSA graph type from Proofs.Autograd.Algebra.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.PNodeData (α Δ : Type) (Γ : List Spec.Shape) (s : Spec.Shape) : Type

Executable node payload for the proof-compiled SSA graph (GraphData).

structure Runtime.Autograd.Compiled.GraphM.Var (s : Spec.Shape) : Type

A typed handle to a value in the growing compiled context.

Var s carries its expected Shape at the type level, while id is the runtime index into the concatenated context Γ ++ ss.

• id : ℕ

  Runtime id of the value inside the concatenated context Γ ++ ss.

  The shape index on Var s is the static guarantee; this numeric id is the executable handle used when constructing Idx proofs for GraphData nodes.

@[implicit_reducible]

instance Runtime.Autograd.Compiled.GraphM.instReprVar {s✝ : Spec.Shape} : Repr (Var s✝)

def Runtime.Autograd.Compiled.GraphM.instReprVar.repr {s✝ : Spec.Shape} : Var s✝ → ℕ → Std.Format

GraphM.arg is correct but a little noisy for examples (you must repeat the index and shape).

VarList + args give a typed variable layer: args returns one Var per entry in Γ, in order, without spelling indices.

inductive Runtime.Autograd.Compiled.GraphM.VarList : List Spec.Shape → Type

Dependent list of typed variables, aligned with a list of shapes.

VarList Γ contains exactly one Var s for each s ∈ Γ, in order.

• nil : VarList []
• cons {s : Spec.Shape} {ss : List Spec.Shape} : Var s → VarList ss → VarList (s :: ss)

def Runtime.Autograd.Compiled.GraphM.VarList.head {s : Spec.Shape} {ss : List Spec.Shape} : VarList (s :: ss) → Var s

First variable in a nonempty VarList.

def Runtime.Autograd.Compiled.GraphM.VarList.tail {s : Spec.Shape} {ss : List Spec.Shape} : VarList (s :: ss) → VarList ss

Tail variables in a nonempty VarList.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.StateWith (α Δ : Type) (Γ : List Spec.Shape) : Type

State for the GraphM builder.

It is a sigma pair of:

• the list of intermediate shapes ss produced so far, and
• the corresponding executable SSA graph payload GraphData α Δ Γ ss.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.State (α : Type) (Γ : List Spec.Shape) : Type

Default GraphM state with no extra environment (Δ := Unit).

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.MWith (α Δ : Type) (Γ : List Spec.Shape) : Type → Type

StateT builder monad for authoring a GraphData program, with explicit environment Δ.

@[reducible, inline]

abbrev Runtime.Autograd.Compiled.GraphM.M (α : Type) (Γ : List Spec.Shape) : Type → Type

Default GraphM builder monad with Δ := Unit.

def Runtime.Autograd.Compiled.GraphM.empty {α : Type} {Γ : List Spec.Shape} : State α Γ

Empty builder state (no intermediate nodes yet).

def Runtime.Autograd.Compiled.GraphM.emptyWith {α Δ : Type} {Γ : List Spec.Shape} : StateWith α Δ Γ

Empty builder state for an explicit environment type Δ.

def Runtime.Autograd.Compiled.GraphM.run {α : Type} {Γ : List Spec.Shape} {β : Type} (m : M α Γ β) : Result (β × State α Γ)

Run a GraphM program from an empty state.

def Runtime.Autograd.Compiled.GraphM.buildGraph {α : Type} {Γ : List Spec.Shape} (m : M α Γ Unit) : Result (State α Γ)

Build a GraphData by running a GraphM program.

This is the usual entry point: write a do-block that constructs the graph using arg, ops, and returns Unit; get back the finalized builder state containing ss and the graph.

def Runtime.Autograd.Compiled.GraphM.ctxLen {Γ : List Spec.Shape} (ss : List Spec.Shape) : ℕ

Length of the current context Γ ++ ss (inputs + intermediates).

def Runtime.Autograd.Compiled.GraphM.mkIdx {_α : Type} [DecidableEq Spec.Shape] {Γ : List Spec.Shape} (ss : List Spec.Shape) {s : Spec.Shape} (v : Var s) : Result (Proofs.Autograd.Algebra.Idx (Γ ++ ss) s)

Convert a Var s into a dependent Idx (Γ ++ ss) s.

This performs bounds checking and a runtime shape check, returning a structured error if the variable points outside the current context or has the wrong shape.

def Runtime.Autograd.Compiled.GraphM.push {α Δ : Type} {Γ ss : List Spec.Shape} {s : Spec.Shape} (g : PGraphData α Δ Γ ss) (node : PNodeData α Δ (Γ ++ ss) s) : MWith α Δ Γ (Var s)

Append a node to the graph state and return a fresh Var pointing to its output.

The returned variable id is Γ.length + ss.length, i.e. it points at the newly appended entry.

inductive Runtime.Autograd.Compiled.GraphM.JvpAvailability : Type

Forward-mode JVP availability for a compiled graph builder op.

• implemented : JvpAvailability

  The op supplies a real forward-mode JVP rule.

• reverseOnly (op : String) : JvpAvailability

  The op supplies reverse-mode VJP only. Forward-mode requests fail loudly.

@[implicit_reducible]

instance Runtime.Autograd.Compiled.GraphM.instReprJvpAvailability : Repr JvpAvailability

def Runtime.Autograd.Compiled.GraphM.instReprJvpAvailability.repr : JvpAvailability → ℕ → Std.Format

def Runtime.Autograd.Compiled.GraphM.instDecidableEqJvpAvailability.decEq (x✝ x✝¹ : JvpAvailability) : Decidable (x✝ = x✝¹)

@[implicit_reducible]

instance Runtime.Autograd.Compiled.GraphM.instDecidableEqJvpAvailability : DecidableEq JvpAvailability

def Runtime.Autograd.Compiled.GraphM.reverseOnlyJvpOps : List String

Compiled ops that provide VJP for training but no forward-mode JVP rule.

Keeping the list executable gives callers a stable preflight hook instead of discovering the gap only after a directional-derivative run reaches the node. The list is intentionally empty when all compiled builder ops have concrete JVP rules.

def Runtime.Autograd.Compiled.GraphM.jvpAvailability (op : String) : JvpAvailability

Return the JVP status for a named compiled op.

def Runtime.Autograd.Compiled.GraphM.reverseOnlyJvpMessage (op : String) : String

Diagnostic message for reverse-only compiled ops.

def Runtime.Autograd.Compiled.GraphM.unsupportedJvp {α : Type} [Zero α] {s : Spec.Shape} (op : String) : Spec.Tensor α s

Fail-fast marker for compiled nodes whose forward-mode JVP rule is intentionally absent.

Returning a zero tangent here would silently corrupt forward-mode autodiff. Reverse-mode users are unaffected because these nodes still provide real vjp implementations. Forward-mode callers get a loud error, and reverseOnlyJvpOps provides a preflight list for tools that want to reject such graphs before running a JVP.

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

Reference an input variable from the initial context Γ.

This checks that the provided index is within bounds and that the requested shape matches the shape at that position in Γ.

PyTorch comparison: this is like naming a graph input tensor in a traced graph.

def Runtime.Autograd.Compiled.GraphM.argsAux (Γ : List Spec.Shape) : ℕ → VarList Γ

Pure helper to build VarList Γ starting at a given id offset.

def Runtime.Autograd.Compiled.GraphM.args {α Δ : Type} {Γ : List Spec.Shape} : MWith α Δ Γ (VarList Γ)

Return one Var per entry of Γ, in order.

This is the canonical argument environment for a graph with input context Γ.

def Runtime.Autograd.Compiled.GraphM.const {α Δ : Type} [Zero α] {Γ : List Spec.Shape} {s : Spec.Shape} (t : Spec.Tensor α s) : MWith α Δ Γ (Var s)

Embed a constant tensor as a node in the compiled graph.

This node has no input dependencies (vjp = 0, jvp = 0), i.e. it is treated as a constant with respect to the graph inputs.

PyTorch comparison: a constant literal captured into a traced/compiled graph.

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

Deterministic U[0,1) tensor generator (seeded, pure).

def Runtime.Autograd.Compiled.GraphM.bernoulliMask {α : Type} [Context α] [DecidableEq Spec.Shape] {Δ : Type} {Γ : List Spec.Shape} {s : Spec.Shape} (keepProb : Var Spec.Shape.scalar) (seed : ℕ) : MWith α Δ Γ (Var s)

Deterministic {0,1} mask generator (seeded, pure).

Note: for differentiation purposes, this node is treated as a stop-gradient op: jvp = 0 and vjp = 0 for all inputs (including keepProb). This matches the intended use in dropout where the probability is a hyperparameter (not differentiated), while keeping execution deterministic in the .compiled backend.

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

Stop-gradient boundary.

Forward semantics: identity (detach(x) = x). Backward semantics: no gradient flows to x (treated as constant w.r.t. the graph inputs).