FX

Active contributors: jamesr66a, ezyang, jansel

Purpose

torch.fx is a Python toolkit for symbolic tracing and transformation of PyTorch programs. It defines a small functional IR (Graph of Nodes), a Tracer that produces graphs from Python code by symbolic execution, and a GraphModule that turns a graph back into runnable Python.

FX is the intermediate representation shared by Dynamo, AOT Autograd, Inductor, the export pipeline, and most graph-rewriting passes. A surprising amount of PyTorch's compiler stack is "an FX graph going through one transformation after another".

Directory layout

Path	Contents
torch/fx/	Core FX
torch/fx/graph.py	Graph, Node, codegen back to Python
torch/fx/graph_module.py	GraphModule
torch/fx/symbolic_trace.py	Default tracer (proxy-based)
torch/fx/proxy.py	Proxy, TraceTransformer
torch/fx/passes/	Reusable graph passes (shape prop, dialect splits, partitioning)
torch/fx/experimental/	Symbolic shapes, proxy_tensor, recording, accelerator partitioner
torch/fx/experimental/symbolic_shapes.py	SymInt/ShapeEnv, dynamic shape constraints
torch/fx/experimental/proxy_tensor.py	The make_fx tracer used by AOT Autograd
torch/csrc/fx/	A small C++ piece for proxy ops

Key abstractions

Type	File	Purpose
Graph	torch/fx/graph.py	A doubly-linked list of Nodes
Node	torch/fx/graph.py	An operation: op ∈ {placeholder, get_attr, call_function, call_module, call_method, output}
GraphModule	torch/fx/graph_module.py	nn.Module whose forward is generated from a Graph
Tracer	torch/fx/symbolic_trace.py	Default symbolic tracer
Proxy	torch/fx/proxy.py	Symbolic value that records ops as it's used
Interpreter	torch/fx/interpreter.py	Walks a Graph, calling ops with concrete values
Transformer	torch/fx/interpreter.py	Walks a Graph and emits a new one
make_fx	torch/fx/experimental/proxy_tensor.py	The Dispatcher-mode tracer used everywhere by torch.compile

How it works

IR shape

A Graph is a list of Nodes. Each Node has:

• op — the kind: placeholder (input), get_attr (module attribute lookup), call_function, call_method, call_module, output.
• target — what to call (a Python callable for call_function, an attribute name for the others).
• args and kwargs — references to other Nodes or Python literals.

The Graph generates Python source on demand (graph.python_code(...)) which the GraphModule wraps as its forward. Pretty-printing is the same generator with simpler formatting.

Tracing modes

PyTorch has two primary tracers that both produce FX graphs:

1. Default tracer (symbolic_trace) — proxy-based. Each input becomes a Proxy; arithmetic operations on Proxys record FX nodes. Works for code that operates on tensor-like proxies with no concrete-Python control flow on tensor values.

2. make_fx (proxy-tensor mode) — the modern, more general tracer used by AOT Autograd, torch.export, torch.compile. Pushes a ProxyTorchDispatchMode and a FakeTensorMode; every dispatcher call is captured at the dispatch level rather than at the Python-call level. This lets it trace through arbitrary ATen op composition, including in-place ops once functionalize has stripped them.

Dynamo doesn't use either tracer — it traces bytecode. But its output is still an FX graph.

Symbolic shapes

torch/fx/experimental/symbolic_shapes.py provides ShapeEnv, the type that tracks the symbolic dimensions used by Dynamo. Tensors flowing through tracing carry SymInt sizes; the env accumulates equality and inequality constraints, propagates them to guards, and is responsible for things like "size 0 vs. size 1 vs. arbitrary" specialization.

This is also where the mark_dynamic and mark_static API lives.

Passes

The torch/fx/passes/ directory has reusable transformations:

• ShapeProp — propagate concrete or fake shapes through a graph.
• split_module — split a graph at named cut points (used by pipelining).
• tools_common — graph splitter / partitioner utilities used by inductor and lazy.
• dialect/ — common, ATen, and Inductor dialects.
• runtime_assert.py — insert runtime asserts for dynamic shape constraints.

User passes typically subclass Transformer or just walk graph.nodes and rewrite in place.

GraphModule lifecycle

graph LR
    User[Python fn / nn.Module] -->|symbolic_trace or make_fx| Graph[Graph]
    Graph -->|GraphModule| GM[GraphModule]
    GM -->|graph.python_code| Src[Python source]
    Src -->|exec| Forward[forward()]
    GM -->|recompile after edits| Forward

GraphModule.recompile() regenerates the Python source after a Graph edit; GraphModule.print_readable() is the standard way to inspect them.

Integration points

• Dynamo emits FX graphs. See Dynamo.
• AOT Autograd uses make_fx to build joint forward+backward graphs. See AOT Autograd.
• Inductor consumes FX graphs and lowers to its own IR. See Inductor.
• Quantization (FX mode) runs as a sequence of FX passes. See Quantization.
• torch.export produces an ExportedProgram whose .graph_module is FX-based. See torch/export/.
• ONNX export (the modern dynamo-onnx exporter) consumes FX. See ONNX.
• Pipelining uses split_module to cut graphs across pipeline stages.

Entry points for modification

• A new graph pass → drop a function under torch/fx/passes/ or in your own module; receive a GraphModule and return one.
• A new node printer/inspector → see Graph.python_code.
• A new tracer → subclass Tracer or use make_fx with custom dispatch modes.
• For dynamic shape work, the file you'll spend time in is torch/fx/experimental/symbolic_shapes.py.

Key source files

File	Purpose
torch/fx/graph.py	Graph and Node
torch/fx/graph_module.py	GraphModule
torch/fx/symbolic_trace.py	Default proxy-based tracer
torch/fx/proxy.py	Proxy and proxy ops
torch/fx/experimental/proxy_tensor.py	make_fx (dispatcher-mode tracer)
torch/fx/experimental/symbolic_shapes.py	SymInt, ShapeEnv, dynamic shapes
torch/fx/passes/shape_prop.py	Shape propagation
torch/fx/passes/split_module.py	Splitting
torch/fx/interpreter.py	Interpreter and Transformer