PyTorch graph capture #

This module generates a Python script that captures an nn.Module and writes a shared tensor-graph artifact. The Lean importer checks it before evaluation:

PyTorch nn.Module
  --torch.export / FX capture-->
TorchLean graph JSON (`torchlean.ir.v2`)
  --Import.IR.parseProgram-->
NN.IR.Program

PyTorch handles capture and shape propagation. The generated script translates supported operations into NN.IR.OpKind and reports an error for unsupported operations or configurations. The importer then checks the graph's edges, shapes, and operation attributes.

Tensor entries also retain their dtypes. Binary constants and linear parameters carry complete interchange words for Import.IR.parseProgram and its typed payload loader. Conversions and pointwise tensor promotions become explicit cast nodes, two of them when an integer, Boolean or binary64 value becomes binary16 or bfloat16, which c10::convert makes from binary32. Pointwise broadcasting is explicit too. Metadata assertions and assumptions needed by a lowering are retained as capture_constraints. For example, partial top-k routing assumes finite scores and a strict gap at the cutoff; capturing one example does not establish those conditions for every input. The CPU reference interpreter does not validate these assumptions or physical device and layout constraints.

example may be one tensor or a tuple of positional tensors. Argument order comes from torch.export's graph signature, not from parameter placeholder order. Unused user arguments remain in the interface; dead parameter values and dead intermediate nodes may be omitted. The artifact lists every input in input_ids, in argument order. Returned containers are flattened to tensor leaves; non-tensor leaves are rejected.

Operator matching uses exact callable or ATen identities. For example, torch.log can become an IR .log node, while torch.log1p needs a different formula and is rejected. Similar names and matching output shapes are insufficient to identify the computation.

Pointwise operators from the core ATen vocabulary (neg, rsqrt, erf, gelu, div, pow and clamp_min with a scalar bound) become IR pointwise nodes in the arithmetic of the vectorized loops of PyTorch's CPU kernels. Operands are converted as PyTorch converts them, Python scalars included, and binary16 and bfloat16 values are widened to binary32 where those loops compute in binary32. A Python scalar operand of add, sub or mul becomes a constant in the same way. For bfloat16 and binary16, a scalar exponent of pow is rounded to the tensor's dtype before PyTorch's dispatch compares it. Exponent 3 is x * x * x, not a rounded cube. The frontend-lowering section of NN/IR/README.md states the rules and the elements PyTorch computes differently.

The script is assembled from one definition per Python section so each piece can be read and changed on its own. Every "kind" string it writes comes from NN.Runtime.PyTorch.Wire, the same table the Lean importer parses with.

References:

torch.export: https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/export.html
torch.fx: https://docs.pytorch.org/docs/stable/fx.html

PyTorch graph capture #

Python sections #