IR Graph

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

Today it is used as the shared target for:

• TorchLean → verifier compilation (NN/Verification/TorchLean/Compile.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).

Longer-term, the intent is to use the same IR as a bridge target for:

• spec-level graphs (compile a model spec to an IR graph),
• runtime autograd traces (reify a runtime tape/graph into the same IR),
• verifiers (IBP/CROWN/affine passes) and export tooling.

What this file commits to is the graph structure (ops, dependencies, and shapes). Parameter payloads (weights/bias/const values) 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).

If you are coming from PyTorch: the mental model is similar to a PyTorch FX graph or TorchScript IR: nodes are ops, edges are “data dependencies”, and execution is in topological order. The difference is that TorchLean attaches explicit shape metadata at every node, since our verification and proof tooling needs shape information to be first-class.

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: in most builders, node.id is expected to equal its index in Graph.nodes. (E.g. the TorchLean compiler 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.conv2d) typically use external parameter stores keyed by node id; in those cases the node’s parents list 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.HardMask : Type

A row-major Boolean mask carried by an IR operation.

The payload records its logical tensor shape separately from the flat array so graph validation can reject malformed serialized or programmatically constructed masks before evaluation. true means that the corresponding entry is allowed.

• shape : Spec.Shape
• allowed : Array Bool

def NN.IR.instReprHardMask.repr : HardMask → ℕ → Std.Format

@[implicit_reducible]

instance NN.IR.instReprHardMask : Repr HardMask

inductive NN.IR.OpKind : Type

Operation kinds in an op-tagged computation graph.

• input : OpKind
• const (valueShape : Spec.Shape) : OpKind
• permute (perm : List ℕ) : OpKind
• detach : OpKind
• randUniform (seed : ℕ) : OpKind
• bernoulliMask (seed : ℕ) : OpKind
• add : OpKind
• sub : OpKind
• mul_elem : OpKind
• abs : OpKind
• sqrt : OpKind
• inv : OpKind
• maxElem : OpKind
• minElem : OpKind
• maxPool2d (kH kW stride : ℕ) : OpKind
• maxPool2dPad (kH kW stride padding : ℕ) : OpKind
• avgPool2d (kH kW stride : ℕ) : OpKind
• avgPool2dPad (kH kW stride padding : ℕ) : OpKind
• broadcastTo (s₁ s₂ : Spec.Shape) : OpKind
• reduceSum (axis : ℕ) : OpKind
• reduceMean (axis : ℕ) : OpKind
• sum : OpKind
• matmul : OpKind
• linear : OpKind
• conv2d (inC outC kH kW stride padding : ℕ) : OpKind
• batchNorm2dNchwEval (channels : ℕ) : OpKind
• relu : OpKind
• tanh : OpKind
• sigmoid : OpKind
• exp : OpKind
• log : OpKind
• sin : OpKind
• cos : OpKind
• softmax (axis : ℕ) : OpKind
• hardMaskedSoftmax (mask : HardMask) : OpKind
• layernorm (axis : ℕ) : OpKind
• reshape (inShape outShape : Spec.Shape) : OpKind
• flatten (s : Spec.Shape) : OpKind
• concat (axis : ℕ) : OpKind
• swap_first_two : OpKind
• transpose3dLastTwo : OpKind
• mseLoss : OpKind

def NN.IR.instReprOpKind.repr : OpKind → ℕ → Std.Format

@[implicit_reducible]

instance NN.IR.instReprOpKind : Repr OpKind

def NN.IR.OpKind.minParents : OpKind → ℕ

The minimum number of parent nodes expected by an OpKind.

This is a structural convention only.

For example, linear has arity 1 here because weights/biases are typically stored externally and keyed by the node id; the only data dependency in the graph itself is the activation input.

def NN.IR.OpKind.maxParents? : OpKind → Option ℕ

An optional maximum number of parent nodes expected by an OpKind.

For concat, the verifier permits an arbitrary number of inputs (at least 2), so this returns none.

def NN.IR.OpKind.tag : OpKind → String

A short tag for error messages and debugging output.

def NN.IR.OpKind.describe : OpKind → String

Human-facing operation description including operation-local parameters.

tag is short and stable for grouping/log filtering. describe is for diagnostics: it prints axes, shapes, seeds, and convolution/pooling metadata so malformed graph dumps are useful without cross-referencing the original builder.

structure NN.IR.Node : Type

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

• id : ℕ

  Node id. By convention this is also the node's index in Graph.nodes.

• parents : List ℕ

  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.

@[implicit_reducible]

instance NN.IR.instReprNode : Repr Node

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

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).

structure NN.IR.Graph : Type

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

• nodes : Array Node

  nodes.

@[implicit_reducible]

instance NN.IR.instReprGraph : Repr Graph

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

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.

def NN.IR.Graph.outShape? (g : Graph) (id : ℕ) : Option Spec.Shape

Safe outShape lookup by id.

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.

@[implicit_reducible]

instance NN.IR.instInhabitedNode : Inhabited Node

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