Shape Inference

Shape inference and consistency checking for NN.IR.Graph.

NN.IR.Node stores an outShape field because many consumers want shape metadata to be available without re-running inference (pretty printers, exporters, verifiers, etc.).

This module provides an independent shape inference/checking procedure that recomputes the expected output shape of each node from:

• the node's OpKind payload (when present), and
• the parent nodes' output shapes.

For parameterized ops whose output shape depends on external parameters (notably OpKind.linear), we treat the node's declared outShape as an input to the checker and validate the local contracts we can check (e.g. a linear layer preserves all leading input dimensions).

Graph.checkShapes uses these rules directly. Adding a new OpKind should extend this match before the semantics, export, and verification passes rely on its shape contract, and should route any nontrivial shape arithmetic through NN.IR.OpContracts so that NN.IR.Semantics computes the same shape; NN.IR.ShapeSoundness proves that agreement operation by operation.

PyTorch analogy:

• nodeOutShape corresponds to shape propagation used when validating an FX graph.
• Where the true output shape depends on parameters, this module performs contract checking rather than attempting to read those parameters.

References / related systems:

• PyTorch FX (graph representation): https://pytorch.org/docs/stable/fx.html
• ONNX shape inference: https://onnx.ai/onnx/shape_inference.html

Node-local inference

Most IR ops are “shape transparent” (elementwise, permute, etc.). A few need special handling:

• matmul preserves an arbitrary shared leading shape around its final matrix axes,
• concat needs to merge multiple parents along an axis,
• pooling and convolution use centralized rank-polymorphic spatial arithmetic from OpContracts.

def NN.IR.Infer.expectUnaryParent (tag : String) (parents : Array Spec.Shape) : Except String Spec.Shape

Read the sole parent shape of a unary operation.

def NN.IR.Infer.expectBinaryParents (tag : String) (parents : Array Spec.Shape) : Except String (Spec.Shape × Spec.Shape)

Read both parent shapes of a binary operation.

def NN.IR.Infer.nodeOutShape (n : Node) (parentShapes : Array Spec.Shape) : Except String Spec.Shape

Infer the output shape of a node from its kind + parent shapes.

This function is used by Graph.checkShapes below.

def NN.IR.Graph.lookupParentShapes (inferred : Array Spec.Shape) : List ℕ → Option (List Spec.Shape)

Look up the already inferred shapes of a node's parents.

Every parent id must index a shape that has already been inferred. On a well-formed graph the topological-order check guarantees this, and the lookup reports a readable error instead of indexing out of bounds if a caller skips that check.

@[irreducible]

def NN.IR.Graph.inferShapesFrom (g : Graph) (i : ℕ) (inferred : Array Spec.Shape) : Except String (Array Spec.Shape)

Infer and check the declared output shapes of nodes i, i+1, ..., extending the table inferred of shapes already inferred for the nodes below i.

The recursion is structural on the remaining node count so proofs can follow it in lockstep with Graph.denoteAllFrom; see NN.IR.ShapeSoundness.

def NN.IR.Graph.inferShapes (g : Graph) : Except String (Array Spec.Shape)

Infer every node's output shape (in topo/id order) after checking structural well-formedness, and check that each Node.outShape matches. Returns the inferred shapes, one per node.

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

Infer shapes for every node (in topo/id order) and check that Node.outShape matches.

This is meant as a lowering/backend consistency check and as a clean IR invariant for the docs: well-formed graphs have self-consistent declared shapes. NN.IR.ShapeSoundness proves that on a graph accepted here the reference semantics computes exactly the declared shapes.