IR Graph

NN.IR.Graph is TorchLean’s canonical op-tagged DAG IR.

Today it is used as the shared target for:

• TorchLean to verifier lowering (NN/Verification/Builtin/Lowering.lean),
• bound-propagation / verification tooling (CROWN/LiRPA) (NN/MLTheory/CROWN/Graph.lean),
• IR → PyTorch emission (NN/Runtime/PyTorch/Export/IRPyTorch.lean),
• compact example graphs (e.g. NN/Examples/DeepDives/GraphSpec/Tutorial.lean).

NN.IR.Operator defines the operations and their static attributes. This file adds node identities, dependencies, and declared output shapes. Parameter payloads (weights, biases, and constants) live in backend-specific stores keyed by node id. This split keeps one graph format usable across:

• verification (where parameters often carry additional metadata like bounds or perturbation sets),
• export (where parameters may be emitted as PyTorch nn.Parameters or ONNX initializers),
• and runtime execution/tracing (where parameters may already live in a separate module state).

Like a PyTorch FX graph or TorchScript IR, nodes are operations, edges are data dependencies, and execution follows topological order. TorchLean additionally attaches explicit shape metadata to every node for verification and proofs.

References / related systems:

• PyTorch FX docs: https://pytorch.org/docs/stable/fx.html
• TorchScript overview: https://pytorch.org/docs/stable/jit.html
• ONNX (graph + initializers as separate parameter store): https://onnx.ai/

Conventions (important)

• Topo order: a node only references parents with smaller ids.
• Id discipline: checked graphs require node.id to equal its index in Graph.nodes. TorchLean lowering uses freshId := nodes.size and then appends.
• External parameters:
  • OpKind.const stores its valueShape here, but the constant value is stored externally (e.g. in a verifier ParamStore keyed by node id).
  • Some ops (notably OpKind.linear and OpKind.conv) typically use external parameter stores keyed by node id; in those cases the node’s parents array only contains the runtime inputs (e.g. the activation input x), not the weights/bias tensors.

This file does not implement evaluation or shape inference. Those live in:

• NN/IR/Semantics.lean (evaluation semantics for a chosen scalar backend),
• NN/IR/Infer.lean / NN/IR/Check.lean (shape inference/checking utilities),
• and backend-specific passes (verification/export) that interpret OpKind in their own setting.

structure NN.IR.Node : Type

Node in the graph. Edges are implicit via parent indices.

• id : ℕ

  Node id. Structural validation requires this to equal the index in Graph.nodes.

• parents : Array ℕ

  Parent node ids, i.e. data dependencies. Each parent must be smaller than id.

• kind : OpKind

  Operation tag and any operation-local metadata.

• outShape : Spec.Shape

  Declared output shape. NN.IR.Infer can recompute/check this from parents.

def NN.IR.instReprNode.repr : Node → ℕ → Std.Format

@[instance_reducible]

instance NN.IR.instReprNode : Repr Node

def NN.IR.Node.hasValidArity (n : Node) : Bool

Check the basic parent-count convention for this node kind.

def NN.IR.Node.parentsBelow (n : Node) : Bool

Check that every parent id is strictly smaller than this node id (topological order).

This is the single most important invariant for the IR:

• it guarantees acyclicity,
• it makes evaluation/inference a simple left-to-right pass,
• and it makes backends predictable (no hidden recursion or “graph rewriting during execution”).

def NN.IR.Node.summary (n : Node) : String

Render a compact, user-facing summary (useful in error messages).

def NN.IR.unaryParent? (parents : Array ℕ) : Option ℕ

Return the sole parent id when an IR node has unary arity.

def NN.IR.binaryParents? (parents : Array ℕ) : Option (ℕ × ℕ)

Return both parent ids when an IR node has binary arity.

theorem NN.IR.mem_of_unaryParent?_eq_some {parents : Array ℕ} {parent : ℕ} (h : unaryParent? parents = some parent) : parent ∈ parents

The parent returned by the unary decoder belongs to the source array.

theorem NN.IR.fst_mem_of_binaryParents?_eq_some {parents : Array ℕ} {left right : ℕ} (h : binaryParents? parents = some (left, right)) : left ∈ parents

The first parent returned by the binary decoder belongs to the source array.

theorem NN.IR.snd_mem_of_binaryParents?_eq_some {parents : Array ℕ} {left right : ℕ} (h : binaryParents? parents = some (left, right)) : right ∈ parents

The second parent returned by the binary decoder belongs to the source array.

structure NN.IR.Graph : Type

Entire graph as an array of nodes. Parents must have smaller ids (topo order).

• nodes : Array Node

  The nodes, in topological order: every parent id is strictly smaller than the index of the node referencing it. Evaluation is then a single left-to-right pass with no scheduling step.

def NN.IR.instReprGraph.repr : Graph → ℕ → Std.Format

@[instance_reducible]

instance NN.IR.instReprGraph : Repr Graph

def NN.IR.Graph.size (g : Graph) : ℕ

Number of nodes in the graph.

def NN.IR.Graph.getNode? (g : Graph) (id : ℕ) : Option Node

Safe node lookup by id (treating ids as array indices).

def NN.IR.Graph.getNode (g : Graph) (id : ℕ) : Except String Node

Total node lookup that enforces the common "id discipline" invariant $\mathrm{nodes}[i].\mathrm{id}=i$.

This is convenient for backends that treat node ids as array indices (verifiers, exporters, pretty printers). The error message is meant to point to a builder bug rather than a user error.

theorem NN.IR.Graph.getNode_id_eq {g : Graph} {id : ℕ} {node : Node} (h : g.getNode id = Except.ok node) : node.id = id

A successful checked lookup returns a node whose stored id is the requested array index.

def NN.IR.Graph.arityError (n : Node) : String

Explain why Node.hasValidArity failed.

This returns a human-facing message rather than structured data; callers use it for diagnostics.

def NN.IR.Graph.wellFormed (g : Graph) : Bool

Basic well-formedness check used by verifier code paths.

This checks:

• node ids match array indices (common construction invariant),
• each node respects its op arity convention, and
• all parent ids are strictly smaller than the node id (topological order).

We keep this as a boolean predicate because some passes want a fast “yes/no” filter. If you need a human-facing error, use checkWellFormed.

def NN.IR.Graph.checkWellFormed (g : Graph) : Except String Unit

Like wellFormed, but returns a helpful error message on failure.

This is useful when you want a clean user error rather than a silent false.

@[instance_reducible]

instance NN.IR.instInhabitedNode : Inhabited Node

Default node used only to satisfy generic container APIs; real graphs should not rely on it.