Skip to content

fx— Symbolic Model Transformation

Overview

torch.fx traces your module into a graph of Nodes and lets you rewrite it— fuse ops, fold constants, quantize, inject hooks, or export. It's the engine under quantization tooling, torch.compile's front-end, and many research projects.

  • symbolic_trace(module) → GraphModule with graph, code(), and full IR.
  • Graph is a DAG of call_function / call_module / call_method / get_attr / placeholder nodes.
  • Write transformations by walking nodes and mutating graph.
  • GraphModule re-compiles to a runnable Python module.

fx tracing executes the module's forward with Proxys, so control flow must be representable: data-dependent branches break naive tracing (fx has Tracer guards).

Trace & Inspect

import torch, torch.nn as nn
from torch.fx import symbolic_trace

class Net(nn.Module):
 def __init__(self):
 super().__init__()
 self.fc1 = nn.Linear(8, 16)
 self.fc2 = nn.Linear(16, 4)
 def forward(self, x):
 return self.fc2(torch.relu(self.fc1(x)))

gm = symbolic_trace(Net())
print(gm.graph) # IR dump
# placeholder x -> call_module fc1 -> call_function relu -> call_module fc2 -> output
print(gm.code) # readable python

Graph Structure (the IR)

for node in gm.graph.nodes:
 print(node.op, node.name, node.target, list(node.args))
# placeholder x x ()
# call_module fc1 fc1 (x,)
# call_function relu <built-in method relu> (fc1,)
# call_module fc2 fc2 (relu,)
# output output (fc2,)

Transformations— Replacing Ops

def replace_relu_with_hardtanh(gm):
 for node in gm.graph.nodes:
 if node.op == 'call_function' and node.target is torch.relu:
 node.target = lambda z: torch.clamp(z, 0, 1) # new callable
 gm.recompile()
 return gm

gm2 = replace_relu_with_hardtanh(symbolic_trace(Net()))
print("code now uses torch.clamp:", "clamp" in gm2.code)

Fuse consecutive linears (demo pass)

def fuse_consecutive_linear(gm):
 linear_nodes = [n for n in gm.graph.nodes
 if n.op == 'call_module' and 'fc' in n.target]
 # real passes inspect weight shapes; this is the shape of a pass:
 print(f"found {len(linear_nodes)} linear nodes to consider")
 return gm

fuse_consecutive_linear(symbolic_trace(Net()))

-

Writing a Pass That Inserts Nodes

from torch.fx import GraphModule
from torch.fx.graph import Graph

def add_relu_after_first_linear(gm):
 g = gm.graph
 with g.inserting_after(list(g.nodes)[1]): # after placeholders/fc1? insert point
 pass
 return gm

Real quantization passes (e.g., torch.ao.quantization.quantize_fx) are fx passes: they insert quantize/dequantize nodes around linear/conv.

Limitations & Where fx Falls Short

Limitation Workaround
Control flow (if data < 0) torch.export with guards / keep in Python
Python builtins opaque use torch.fx Tracer overrides, or torch.fx.wrap
In-place & aliasing avoid in traced code; use functional forms
Non-tensor args must be traced constants or torch.fx.wrap them

Wrap a python helper

import torch.fx as fx

@fx.wrap
def my_helper(a, b):
 return a.clamp(min=b)

class W(nn.Module):
 def forward(self, x):
 return my_helper(x, 0)
print(list(symbolic_trace(W()).graph.nodes))

fx ↔ Compilation Ecosystem

  • torch.compile(model) uses dynamo (not fx) today, but fx passes remain for research & tooling.
  • torch.ao.quantization.quantize_fx = official fx-based quantization.
  • torch.fx is also the front-end of torch.unittest-style graph tests.
model = torch.ao.quantization.quantize_fx.prepare_qat_fx(
 symbolic_trace(Net()), {'': torch.ao.quantization.default_qat_qconfig}, None)

-

Key Takeaways

  • fx = trace → IR (nodes) → rewrite → recompile (GraphModule).
  • Passes mutate node targets/args and call recompile().
  • Ideal for fusion, folding, and quantization, not for data-dependent control flow.
  • torch.compile/dynamo supersedes fx for everyday acceleration, but fx remains the tooling layer.

-