Lab 01 — FX Graph Passes

Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

FX graph passes are how torch.compile transforms your model before lowering to a backend.

What you build

  • symbolic_trace a model to capture its FX graph
  • Write a graph pass that fuses linear → relu into a single node
  • Write a graph pass that eliminates redundant operations (dead code)
  • Verify output numerical equivalence before and after transformation

Key concepts

ConceptWhat to understand
FX IRThree-address code representation of PyTorch operations
Node kindsplaceholder, get_attr, call_function, call_module, output
Graph rewritingInsert/erase nodes; update users; maintain invariants
Operator fusionFused kernels reduce memory bandwidth (key for NPU perf)

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v