Rate this Page
★★★★★

torch.fx#

Created On: Dec 15, 2020 | Last Updated On: Jul 15, 2025

Overview#

FX is a toolkit for developers to use to transformnn.Moduleinstances. FX consists of three main components: asymbolic tracer,anintermediate representation, andPython code generation. Ademonstration of these components in action:

importtorch# Simple module for demonstrationclassMyModule(torch.nn.Module):def__init__(self)->None:super().__init__()self.param=torch.nn.Parameter(torch.rand(3,4))self.linear=torch.nn.Linear(4,5)defforward(self,x):returnself.linear(x+self.param).clamp(min=0.0,max=1.0)module=MyModule()fromtorch.fximportsymbolic_trace# Symbolic tracing frontend - captures the semantics of the modulesymbolic_traced:torch.fx.GraphModule=symbolic_trace(module)# High-level intermediate representation (IR) - Graph representationprint(symbolic_traced.graph)"""graph():    %x : [num_users=1] = placeholder[target=x]    %param : [num_users=1] = get_attr[target=param]    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})    %linear : [num_users=1] = call_module[target=linear](args = (%add,), kwargs = {})    %clamp : [num_users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})    return clamp"""# Code generation - valid Python codeprint(symbolic_traced.code)"""def forward(self, x):    param = self.param    add = x + param;  x = param = None    linear = self.linear(add);  add = None    clamp = linear.clamp(min = 0.0, max = 1.0);  linear = None    return clamp"""

Thesymbolic tracer performs “symbolic execution” of the Pythoncode. It feeds fake values, called Proxies, through the code. Operationson these Proxies are recorded. More information about symbolic tracingcan be found in thesymbolic_trace() andTracerdocumentation.

Theintermediate representation is the container for the operationsthat were recorded during symbolic tracing. It consists of a list ofNodes that represent function inputs, callsites (to functions, methods,ortorch.nn.Module instances), and return values. More informationabout the IR can be found in the documentation forGraph. TheIR is the format on which transformations are applied.

Python code generation is what makes FX a Python-to-Python (orModule-to-Module) transformation toolkit. For each Graph IR, we cancreate valid Python code matching the Graph’s semantics. Thisfunctionality is wrapped up inGraphModule, which is atorch.nn.Module instance that holds aGraph as well as aforward method generated from the Graph.

Taken together, this pipeline of components (symbolic tracing ->intermediate representation -> transforms -> Python code generation)constitutes the Python-to-Python transformation pipeline of FX. Inaddition, these components can be used separately. For example,symbolic tracing can be used in isolation to capture a form ofthe code for analysis (and not transformation) purposes. Codegeneration can be used for programmatically generating models, forexample from a config file. There are many uses for FX!

Several example transformations can be found at theexamplesrepository.

Writing Transformations#

What is an FX transform? Essentially, it’s a function that looks like this.

importtorchimporttorch.fxdeftransform(m:nn.Module,tracer_class:type=torch.fx.Tracer)->torch.nn.Module:# Step 1: Acquire a Graph representing the code in `m`# NOTE: torch.fx.symbolic_trace is a wrapper around a call to# fx.Tracer.trace and constructing a GraphModule. We'll# split that out in our transform to allow the caller to# customize tracing behavior.graph:torch.fx.Graph=tracer_class().trace(m)# Step 2: Modify this Graph or create a new onegraph=...# Step 3: Construct a Module to returnreturntorch.fx.GraphModule(m,graph)

Your transform will take in atorch.nn.Module, acquire aGraphfrom it, do some modifications, and return a newtorch.nn.Module. You should think of thetorch.nn.Module that your FXtransform returns as identical to a regulartorch.nn.Module – you can pass it to anotherFX transform, or you can run it. Ensuring that the inputs and outputs of your FX transform are atorch.nn.Module will allow for composability.

Note

It is also possible to modify an existingGraphModule instead ofcreating a new one, like so:

importtorchimporttorch.fxdeftransform(m:nn.Module)->nn.Module:gm:torch.fx.GraphModule=torch.fx.symbolic_trace(m)# Modify gm.graph# <...># Recompile the forward() method of `gm` from its Graphgm.recompile()returngm

Note that you MUST callGraphModule.recompile() to bring the generatedforward() method on theGraphModule in sync with the modifiedGraph.

Given that you’ve passed in atorch.nn.Module that has been traced into aGraph, there are now two primary approaches you can take to building a newGraph.

A Quick Primer on Graphs#

Full treatment of the semantics of graphs can be found in theGraphdocumentation, but we are going to cover the basics here. AGraph isa data structure that represents a method on aGraphModule. Theinformation that this requires is:

  • What are the inputs to the method?

  • What are the operations that run inside the method?

  • What is the output (i.e. return) value from the method?

All three of these concepts are represented withNode instances.Let’s see what we mean by that with a short example:

importtorchimporttorch.fxclassMyModule(torch.nn.Module):def__init__(self):super().__init__()self.param=torch.nn.Parameter(torch.rand(3,4))self.linear=torch.nn.Linear(4,5)defforward(self,x):returntorch.topk(torch.sum(self.linear(x+self.linear.weight).relu(),dim=-1),3)m=MyModule()gm=torch.fx.symbolic_trace(m)gm.graph.print_tabular()

Here we define a moduleMyModule for demonstration purposes, instantiate it,symbolically trace it, then call theGraph.print_tabular() method to printout a table showing the nodes of thisGraph:

opcode

name

target

args

kwargs

placeholder

x

x

()

{}

get_attr

linear_weight

linear.weight

()

{}

call_function

add_1

(x, linear_weight)

{}

call_module

linear_1

linear

(add_1,)

{}

call_method

relu_1

relu

(linear_1,)

{}

call_function

sum_1

<built-in method sum …>

(relu_1,)

{‘dim’: -1}

call_function

topk_1

<built-in method topk …>

(sum_1, 3)

{}

output

output

output

(topk_1,)

{}

We can use this information to answer the questions we posed above.

  • What are the inputs to the method? In FX, method inputs are specifiedvia specialplaceholder nodes. In this case, we have a singleplaceholder node with atarget ofx, meaning we havea single (non-self) argument named x.

  • What are the operations within the method? Theget_attr,call_function,call_module, andcall_method nodesrepresent the operations in the method. A full treatment ofthe semantics of all of these can be found in theNodedocumentation.

  • What is the return value of the method? The return value in aGraph is specified by a specialoutput node.

Given that we now know the basics of how code is represented inFX, we can now explore how we would edit aGraph.

Graph Manipulation#

Direct Graph Manipulation#

One approach to building this newGraph is to directly manipulate your oldone. To aid in this, we can simply take theGraph we obtain from symbolictracing and modify it. For example, let’s say we desire to replacetorch.add() calls withtorch.mul() calls.

importtorchimporttorch.fx# Sample moduleclassM(torch.nn.Module):defforward(self,x,y):returntorch.add(x,y)deftransform(m:torch.nn.Module,tracer_class:type=fx.Tracer)->torch.nn.Module:graph:fx.Graph=tracer_class().trace(m)# FX represents its Graph as an ordered list of# nodes, so we can iterate through them.fornodeingraph.nodes:# Checks if we're calling a function (i.e:# torch.add)ifnode.op=='call_function':# The target attribute is the function# that call_function calls.ifnode.target==torch.add:node.target=torch.mulgraph.lint()# Does some checks to make sure the# Graph is well-formed.returnfx.GraphModule(m,graph)

We can also do more involvedGraph rewrites, such asdeleting or appending nodes. To aid in these transformations,FX has utility functions for transforming the graph that canbe found in theGraph documentation. Anexample of using these APIs to append atorch.relu() callcan be found below.

# Specifies the insertion point. Any nodes added to the# Graph within this scope will be inserted after `node`withtraced.graph.inserting_after(node):# Insert a new `call_function` node calling `torch.relu`new_node=traced.graph.call_function(torch.relu,args=(node,))# We want all places that used the value of `node` to# now use that value after the `relu` call we've added.# We use the `replace_all_uses_with` API to do this.node.replace_all_uses_with(new_node)

For simple transformations that only consist of substitutions, you can alsomake use of thesubgraph rewriter.

Subgraph Rewriting With replace_pattern()#

FX also provides another level of automation on top of direct graph manipulation.Thereplace_pattern() API is essentially a “find/replace” tool for editingGraph\s. It allows you to specify apattern andreplacement functionand it will trace through those functions, find instances of the group of operationsin thepattern graph, and replace those instances with copies of thereplacementgraph. This can help to greatly automate tedious graph manipulation code, which canget unwieldy as the transformations get more complex.

Proxy/Retracing#

Another way of manipulatingGraph\s is by reusing theProxymachinery used in symbolic tracing. For example, let’simagine that we wanted to write a transformation that decomposedPyTorch functions into smaller operations. It would transform everyF.relu(x) call into(x>0)*x. One possibility would be toperform the requisite graph rewriting to insert the comparison andmultiplication after theF.relu, and then clean up the originalF.relu. However, we can automate this process by usingProxyobjects to automatically record operations into theGraph.

To use this method, we write the operations that we want inserted as regularPyTorch code and invoke that code withProxy objects as arguments.TheseProxy objects will capture the operations that are performedon them and append them to theGraph.

# Note that this decomposition rule can be read as regular Pythondefrelu_decomposition(x):return(x>0)*xdecomposition_rules={}decomposition_rules[F.relu]=relu_decompositiondefdecompose(model:torch.nn.Module,tracer_class:type=fx.Tracer)->torch.nn.Module:"""    Decompose `model` into smaller constituent operations.    Currently,this only supports decomposing ReLU into its    mathematical definition: (x > 0) * x    """graph:fx.Graph=tracer_class().trace(model)new_graph=fx.Graph()env={}tracer=torch.fx.proxy.GraphAppendingTracer(new_graph)fornodeingraph.nodes:ifnode.op=='call_function'andnode.targetindecomposition_rules:# By wrapping the arguments with proxies,# we can dispatch to the appropriate# decomposition rule and implicitly add it# to the Graph by symbolically tracing it.proxy_args=[fx.Proxy(env[x.name],tracer)ifisinstance(x,fx.Node)elsexforxinnode.args]output_proxy=decomposition_rules[node.target](*proxy_args)# Operations on `Proxy` always yield new `Proxy`s, and the# return value of our decomposition rule is no exception.# We need to extract the underlying `Node` from the `Proxy`# to use it in subsequent iterations of this transform.new_node=output_proxy.nodeenv[node.name]=new_nodeelse:# Default case: we don't have a decomposition rule for this# node, so just copy the node over into the new graph.new_node=new_graph.node_copy(node,lambdax:env[x.name])env[node.name]=new_nodereturnfx.GraphModule(model,new_graph)

In addition to avoiding explicit graph manipulation, usingProxy\salso allows you to specify your rewrite rules as native Python code.For transformations that require a large amount of rewrite rules(such as vmap or grad), this can often improve readability andmaintainability of the rules. Note that while callingProxy we alsopassed a tracer pointing to the underlying variablegraph. This is done soif in case the operations in graph are n-ary (e.g. add is a binary operator)the call toProxy does not create multiple instances of a graphtracer which can lead to unexpected runtime errors. We recommend this methodof usingProxy especially when the underlying operators can not besafely assumed to be unary.

A worked example of usingProxy\s forGraph manipulationcan be foundhere.

The Interpreter Pattern#

A useful code organizational pattern in FX is to loop over all theNode\sin aGraph and execute them. This can be used for several things includingruntime analysis of values flowing through the graph or transformation of the codevia retracing withProxy\s. For example, suppose we want to run aGraphModule and record thetorch.Tensor shape and dtypeproperties on the nodes as we see them at runtime. That might look like:

importtorchimporttorch.fxfromtorch.fx.nodeimportNodefromtypingimportDictclassShapeProp:"""    Shape propagation. This class takes a `GraphModule`.    Then, its `propagate` method executes the `GraphModule`    node-by-node with the given arguments. As each operation    executes, the ShapeProp class stores away the shape and    element type for the output values of each operation on    the `shape` and `dtype` attributes of the operation's    `Node`.    """def__init__(self,mod):self.mod=modself.graph=mod.graphself.modules=dict(self.mod.named_modules())defpropagate(self,*args):args_iter=iter(args)env:Dict[str,Node]={}defload_arg(a):returntorch.fx.graph.map_arg(a,lambdan:env[n.name])deffetch_attr(target:str):target_atoms=target.split('.')attr_itr=self.modfori,atominenumerate(target_atoms):ifnothasattr(attr_itr,atom):raiseRuntimeError(f"Node referenced nonexistent target{'.'.join(target_atoms[:i])}")attr_itr=getattr(attr_itr,atom)returnattr_itrfornodeinself.graph.nodes:ifnode.op=='placeholder':result=next(args_iter)elifnode.op=='get_attr':result=fetch_attr(node.target)elifnode.op=='call_function':result=node.target(*load_arg(node.args),**load_arg(node.kwargs))elifnode.op=='call_method':self_obj,*args=load_arg(node.args)kwargs=load_arg(node.kwargs)result=getattr(self_obj,node.target)(*args,**kwargs)elifnode.op=='call_module':result=self.modules[node.target](*load_arg(node.args),**load_arg(node.kwargs))# This is the only code specific to shape propagation.# you can delete this `if` branch and this becomes# a generic GraphModule interpreter.ifisinstance(result,torch.Tensor):node.shape=result.shapenode.dtype=result.dtypeenv[node.name]=resultreturnload_arg(self.graph.result)

As you can see, a full interpreter for FX is not that complicatedbut it can be very useful. To ease using this pattern, we providetheInterpreter class, which encompasses the above logicin a way that certain aspects of the interpreter’s execution canbe overridden via method overrides.

In addition to executing operations, we can also generate a newGraph by feedingProxy values through an interpreter.Similarly, we provide theTransformer class to encompassthis pattern.Transformer behaves similarly toInterpreter, but instead of calling therun method toget a concrete output value from the Module, you would call theTransformer.transform() method to return a newGraphModule which was subject to any transformation rulesyou installed as overridden methods.

Examples of the Interpreter Pattern#

Debugging#

Introduction#

Often in the course of authoring transformations, our code will not be quite right.In this case, we may need to do some debugging. The key is to workbackwards: first, check the results of invoking the generated module to prove ordisprove correctness. Then, inspect and debug the generated code. Then, debug theprocess of transformations that led to the generated code.

If you’re not familiar with debuggers, please see the auxiliary sectionAvailable Debuggers.

Common Pitfalls in Transform Authoring#

  • Nondeterministicset iteration order. In Python, theset datatype isunordered. Usingset to contain collections of objects likeNode\ s,for example, can cause unexpected nondeterminism. An example is iteratingover a set ofNode s to insert them into aGraph. Because theset data type is unordered, the ordering of the operations in the outputprogram will be nondeterministic and can change across program invocations.The recommended alternative is to use adict data type, which isinsertion orderedas of Python 3.7 (and as of cPython 3.6). Adict can be used equivalentlyto a set by storing values to be deduplicated in the keys of thedict.

Checking Correctness of Modules#

Because the output of most deep learning modules consists of floatingpointtorch.Tensor instances, checking for equivalence betweenthe results of twotorch.nn.Module is not as straightforwardas doing a simple equality check. To motivate this, let’s use anexample:

importtorchimporttorch.fximporttorchvision.modelsasmodelsdeftransform(m:torch.nn.Module)->torch.nn.Module:gm=torch.fx.symbolic_trace(m)# Imagine we're doing some transforms here# <...>gm.recompile()returngmresnet18=models.resnet18()transformed_resnet18=transform(resnet18)input_image=torch.randn(5,3,224,224)assertresnet18(input_image)==transformed_resnet18(input_image)"""RuntimeError: Boolean value of Tensor with more than one value is ambiguous"""

Here, we’ve tried to check equality of the values of two deep learningmodels with the== equality operator. However, this is not well-
defined both due to the issue of that operator returning a tensorand not a bool, but also because comparison of floating point valuesshould use a margin of error (or epsilon) to account for thenon-commutativity of floating point operations (seehere for moredetails). We can usetorch.allclose() instead, which will giveus an approximate comparison taking into account a relative andabsolute tolerance threshold:

asserttorch.allclose(resnet18(input_image),transformed_resnet18(input_image))

This is the first tool in our toolbox to check if transformed modules arebehaving as we expect compared to a reference implementation.

Debugging the Generated Code#

Because FX generates theforward() function onGraphModule\s, usingtraditional debugging techniques likeprint statements orpdb isnot as straightforward. Luckily, we have several techniques we can usefor debugging the generated code.

Usepdb#

Invokepdb to step into the running program. Although the code thatrepresents theGraph is not in any source file, we can still stepinto it manually usingpdb when the forward pass is invoked.

importtorchimporttorch.fximporttorchvision.modelsasmodelsdefmy_pass(inp:torch.nn.Module,tracer_class:type=fx.Tracer)->torch.nn.Module:graph=tracer_class().trace(inp)# Transformation logic here# <...># Return new Modulereturnfx.GraphModule(inp,graph)my_module=models.resnet18()my_module_transformed=my_pass(my_module)input_value=torch.randn(5,3,224,224)# When this line is executed at runtime, we will be dropped into an# interactive `pdb` prompt. We can use the `step` or `s` command to# step into the execution of the next lineimportpdb;pdb.set_trace()my_module_transformed(input_value)

Print the Generated Code#

If you’d like to run the same code multiple times, then it can bea bit tedious to step to the right code withpdb. In that case, oneapproach is to simply copy-paste the generatedforward pass intoyour code and examine it from there.

# Assume that `traced` is a GraphModule that has undergone some# number of transforms# Copy this code for laterprint(traced)# Print the code generated from symbolic tracing. This outputs:"""def forward(self, y):    x = self.x    add_1 = x + y;  x = y = None    return add_1"""# Subclass the original ModuleclassSubclassM(M):def__init__(self):super().__init__()# Paste the generated `forward` function (the one we printed and# copied above) heredefforward(self,y):x=self.xadd_1=x+y;x=y=Nonereturnadd_1# Create an instance of the original, untraced Module. Then, create an# instance of the Module with the copied `forward` function. We can# now compare the output of both the original and the traced version.pre_trace=M()post_trace=SubclassM()

Use theto_folder Function FromGraphModule#

GraphModule.to_folder() is a method inGraphModule that allowsyou to dump out the generated FX code to a folder. Although copying theforward pass into the code often suffices as inPrint the Generated Code,it may be easier to examine modules and parameters usingto_folder.

m=symbolic_trace(M())m.to_folder("foo","Bar")fromfooimportBary=Bar()

After running the above example, we can then look at the code withinfoo/module.py and modify it as desired (e.g. addingprintstatements or usingpdb) to debug the generated code.

Debugging the Transformation#

Now that we’ve identified that a transformation is creating incorrectcode, it’s time to debug the transformation itself. First, we’ll checktheLimitations of Symbolic Tracing section in the documentation.Once we verify that tracing is working as expected, the goalbecomes figuring out what went wrong during ourGraphModuletransformation. There may be a quick answer inWriting Transformations, but, if not, there are several ways toexamine our traced module:

# Sample ModuleclassM(torch.nn.Module):defforward(self,x,y):returnx+y# Create an instance of `M`m=M()# Symbolically trace an instance of `M` (returns a GraphModule). In# this example, we'll only be discussing how to inspect a# GraphModule, so we aren't showing any sample transforms for the# sake of brevity.traced=symbolic_trace(m)# Print the code produced by tracing the module.print(traced)# The generated `forward` function is:"""def forward(self, x, y):    add = x + y;  x = y = None    return add"""# Print the internal Graph.print(traced.graph)# This print-out returns:"""graph():    %x : [num_users=1] = placeholder[target=x]    %y : [num_users=1] = placeholder[target=y]    %add : [num_users=1] = call_function[target=operator.add](args = (%x, %y), kwargs = {})    return add"""# Print a tabular representation of the internal Graph.traced.graph.print_tabular()# This gives us:"""opcode         name    target                   args    kwargs-------------  ------  -----------------------  ------  --------placeholder    x       x                        ()      {}placeholder    y       y                        ()      {}call_function  add     <built-in function add>  (x, y)  {}output         output  output                   (add,)  {}"""

Using the utility functions above, we can compare our traced Modulebefore and after we’ve applied our transformations. Sometimes, asimple visual comparison is enough to trace down a bug. If it’s stillnot clear what’s going wrong, a debugger likepdb can be a goodnext step.

Going off of the example above, consider the following code:

# Sample user-defined functiondeftransform_graph(module:torch.nn.Module,tracer_class:type=fx.Tracer)->torch.nn.Module:# Get the Graph from our traced Moduleg=tracer_class().trace(module)"""    Transformations on `g` go here    """returnfx.GraphModule(module,g)# Transform the Graphtransformed=transform_graph(traced)# Print the new code after our transforms. Check to see if it was# what we expectedprint(transformed)

Using the above example, let’s say that the call toprint(traced)showed us that there was an error in our transforms. We want to findwhat goes wrong using a debugger. We start apdb session. We can seewhat’s happening during the transform by breaking ontransform_graph(traced), then pressings to “step into” the calltotransform_graph(traced).

We may also have good luck by editing theprint_tabular method to printdifferent attributes of the Nodes in the Graph. (For example, we mightwant to see the Node’sinput_nodes andusers.)

Available Debuggers#

The most common Python debugger ispdb. You can startyour program in “debug mode” withpdb by typingpython-mpdbFILENAME.py into the command line, whereFILENAMEis the name of the file you want to debug. After that, you can use thepdbdebugger commandsto move through your running program stepwise. It’s common to set abreakpoint (bLINE-NUMBER) when you startpdb, then callc torun the program until that point. This prevents you from having to stepthrough each line of execution (usings orn) to get to the partof the code you want to examine. Alternatively, you can writeimportpdb;pdb.set_trace() before the line you want to break at.If you addpdb.set_trace(), your program will automatically startin debug mode when you run it. (In other words, you can just typepythonFILENAME.py into the command line instead ofpython-mpdbFILENAME.py.) Once you’re running your file indebug mode, you can step through the code and examine your program’sinternal state using certain commands. There are many excellenttutorials onpdb online, including RealPython’s“Python Debugging With Pdb”.

IDEs like PyCharm or VSCode usually have a debugger built in. In yourIDE, you can choose to either a) usepdb by pulling up a terminalwindow in your IDE (e.g. View → Terminal in VSCode), or b) use thebuilt-in debugger (usually a graphical wrapper aroundpdb).

Limitations of Symbolic Tracing#

FX uses a system ofsymbolic tracing (a.k.asymbolicexecution)to capture the semantics of programs in a transformable/analyzable form.The system istracing in that it executes the program (really atorch.nn.Module or function) to record operations. It issymbolic in that the data flowing through the program during thisexecution is not real data, but rather symbols (Proxy in FX parlance).

Although symbolic tracing works for most neural net code, it has somelimitations.

Dynamic Control Flow#

The main limitation of symbolic tracing is it does not currently supportdynamic control flow. That is, loops orif statements where thecondition may depend on the input values of the program.

For example, let’s examine the following program:

deffunc_to_trace(x):ifx.sum()>0:returntorch.relu(x)else:returntorch.neg(x)traced=torch.fx.symbolic_trace(func_to_trace)"""    <...>    File "dyn.py", line 6, in func_to_trace    if x.sum() > 0:    File "pytorch/torch/fx/proxy.py", line 155, in __bool__    return self.tracer.to_bool(self)    File "pytorch/torch/fx/proxy.py", line 85, in to_bool    raise TraceError('symbolically traced variables cannot be used as inputs to control flow')torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow"""

The condition to theif statement relies on the value ofx.sum(),which relies on the value ofx, a function input. Sincex can change (i.e. if you pass a new input tensor to the tracedfunction), this isdynamic control flow. The traceback walks back upthrough your code to show you where this situation happens.

Static Control Flow#

On the other hand, so-calledstatic control flow is supported. Staticcontrol flow is loops orif statements whose value cannot changeacross invocations. Typically, in PyTorch programs, this control flowarises for code making decisions about a model’s architecture based onhyper-parameters. As a concrete example:

importtorchimporttorch.fxclassMyModule(torch.nn.Module):def__init__(self,do_activation:bool=False):super().__init__()self.do_activation=do_activationself.linear=torch.nn.Linear(512,512)defforward(self,x):x=self.linear(x)# This if-statement is so-called static control flow.# Its condition does not depend on any input valuesifself.do_activation:x=torch.relu(x)returnxwithout_activation=MyModule(do_activation=False)with_activation=MyModule(do_activation=True)traced_without_activation=torch.fx.symbolic_trace(without_activation)print(traced_without_activation.code)"""def forward(self, x):    linear_1 = self.linear(x);  x = None    return linear_1"""traced_with_activation=torch.fx.symbolic_trace(with_activation)print(traced_with_activation.code)"""import torchdef forward(self, x):    linear_1 = self.linear(x);  x = None    relu_1 = torch.relu(linear_1);  linear_1 = None    return relu_1"""

The if-statementifself.do_activation does not depend on anyfunction inputs, thus it is static.do_activation can be consideredto be a hyper-parameter, and the traces of different instances ofMyModule with different values for that parameter have differentcode. This is a valid pattern that is supported by symbolic tracing.

Many instances of dynamic control flow are semantically static controlflow. These instances can be made to support symbolic tracing byremoving the data dependencies on input values, for example by movingvalues toModule attributes or by binding concrete values to argumentsduring symbolic tracing:

deff(x,flag):ifflag:returnxelse:returnx*2fx.symbolic_trace(f)# Fails!fx.symbolic_trace(f,concrete_args={'flag':True})

In the case of truly dynamic control flow, the sections of the programthat contain this code can be traced as calls to the Method (seeCustomizing Tracing with the Tracer class) or function (seewrap()) rather than tracing through them.

Non-torch Functions#

FX uses__torch_function__ as the mechanism by which it interceptscalls (see thetechnicaloverviewfor more information about this). Some functions, such as builtin Pythonfunctions or those in themath module, are not covered by__torch_function__, but we would still like to capture them insymbolic tracing. For example:

importtorchimporttorch.fxfrommathimportsqrtdefnormalize(x):"""    Normalize `x` by the size of the batch dimension    """returnx/sqrt(len(x))# It's valid Python codenormalize(torch.rand(3,4))traced=torch.fx.symbolic_trace(normalize)"""    <...>    File "sqrt.py", line 9, in normalize    return x / sqrt(len(x))    File "pytorch/torch/fx/proxy.py", line 161, in __len__    raise RuntimeError("'len' is not supported in symbolic tracing by default. If you want "RuntimeError: 'len' is not supported in symbolic tracing by default. If you want this call to be recorded, please call torch.fx.wrap('len') at module scope"""

The error tells us that the built-in functionlen is not supported.We can make it so that functions like this are recorded in the trace asdirect calls using thewrap() API:

torch.fx.wrap('len')torch.fx.wrap('sqrt')traced=torch.fx.symbolic_trace(normalize)print(traced.code)"""import mathdef forward(self, x):    len_1 = len(x)    sqrt_1 = math.sqrt(len_1);  len_1 = None    truediv = x / sqrt_1;  x = sqrt_1 = None    return truediv"""

Customizing Tracing with theTracer class#

TheTracer class is the class that underlies theimplementation ofsymbolic_trace. The behavior of tracing can becustomized by subclassing Tracer, like so:

classMyCustomTracer(torch.fx.Tracer):# Inside here you can override various methods# to customize tracing. See the `Tracer` API# referencepass# Let's use this custom tracer to trace through this moduleclassMyModule(torch.nn.Module):defforward(self,x):returntorch.relu(x)+torch.ones(3,4)mod=MyModule()traced_graph=MyCustomTracer().trace(mod)# trace() returns a Graph. Let's wrap it up in a# GraphModule to make it runnabletraced=torch.fx.GraphModule(mod,traced_graph)

Leaf Modules#

Leaf Modules are the modules that appear as calls in the symbolic tracerather than being traced through. The default set of leaf modules is theset of standardtorch.nn module instances. For example:

classMySpecialSubmodule(torch.nn.Module):defforward(self,x):returntorch.neg(x)classMyModule(torch.nn.Module):def__init__(self):super().__init__()self.linear=torch.nn.Linear(3,4)self.submod=MySpecialSubmodule()defforward(self,x):returnself.submod(self.linear(x))traced=torch.fx.symbolic_trace(MyModule())print(traced.code)# `linear` is preserved as a call, yet `submod` is traced though.# This is because the default set of "Leaf Modules" includes all# standard `torch.nn` modules."""import torchdef forward(self, x):    linear_1 = self.linear(x);  x = None    neg_1 = torch.neg(linear_1);  linear_1 = None    return neg_1"""

The set of leaf modules can be customized by overridingTracer.is_leaf_module().

Miscellanea#

  • Tensor constructors (e.g.torch.zeros,torch.ones,torch.rand,torch.randn,torch.sparse_coo_tensor)are currently not traceable.

    • The deterministic constructors (zeros,ones) can be usedand the value they produce will be embedded in the trace as aconstant. This is only problematic if the arguments to theseconstructors refers to dynamic input sizes. In this case,ones_like orzeros_like may be a viable substitute.

    • Nondeterministic constructors (rand,randn) will have asingle random value embedded in the trace. This is likely not theintended behavior. One workaround is to wraptorch.randn in atorch.fx.wrap function and call that instead.

    @torch.fx.wrapdeftorch_randn(x,shape):returntorch.randn(shape)deff(x):returnx+torch_randn(x,5)fx.symbolic_trace(f)

    • This behavior may be fixed in a future release.

  • Type annotations

    • Python 3-style type annotations (e.g.func(x:torch.Tensor,y:int)->torch.Tensor) are supportedand will be preserved by symbolic tracing.

    • Python 2-style comment type annotations#type:(torch.Tensor,int)->torch.Tensor are not currentlysupported.

    • Annotations on local names within a function are not currentlysupported.

  • Gotcha aroundtraining flag and submodules

    • When using functionals liketorch.nn.functional.dropout, it will be common for the training argument to be passed in asself.training. During FX tracing, this will likely be baked in as a constant value.

    importtorchimporttorch.fxclassDropoutRepro(torch.nn.Module):defforward(self,x):returntorch.nn.functional.dropout(x,training=self.training)traced=torch.fx.symbolic_trace(DropoutRepro())print(traced.code)"""def forward(self, x):    dropout = torch.nn.functional.dropout(x, p = 0.5, training = True, inplace = False);  x = None    return dropout"""traced.eval()x=torch.randn(5,3)torch.testing.assert_close(traced(x),x)"""AssertionError: Tensor-likes are not close!Mismatched elements: 15 / 15 (100.0%)Greatest absolute difference: 1.6207983493804932 at index (0, 2) (up to 1e-05 allowed)Greatest relative difference: 1.0 at index (0, 0) (up to 0.0001 allowed)"""

    • However, when the standardnn.Dropout() submodule is used, the training flag is encapsulated and–because of the preservation of thenn.Module object model–can be changed.

    classDropoutRepro2(torch.nn.Module):def__init__(self):super().__init__()self.drop=torch.nn.Dropout()defforward(self,x):returnself.drop(x)traced=torch.fx.symbolic_trace(DropoutRepro2())print(traced.code)"""def forward(self, x):    drop = self.drop(x);  x = None    return drop"""traced.eval()x=torch.randn(5,3)torch.testing.assert_close(traced(x),x)

  • Because of this difference, consider marking modules that interact with thetraining flag dynamically as leaf modules.

API Reference#

torch.fx.symbolic_trace(root,concrete_args=None)[source]#

Symbolic tracing API

Given annn.Module or function instanceroot, this function will return aGraphModuleconstructed by recording operations seen while tracing throughroot.

concrete_args allows you to partially specialize your function, whether it’s to remove control flow or data structures.

For example:

deff(a,b):ifb==True:returnaelse:returna*2

FX can typically not trace through this due to the presence of controlflow. However, we can useconcrete_args to specialize on the value ofb to trace through this:

f=fx.symbolic_trace(f,concrete_args={"b":False})assertf(3,False)==6

Note that although you can still pass in different values ofb, they will be ignored.

We can also useconcrete_args to eliminate data-structure handling fromour function. This will use pytrees to flatten your input. To avoidoverspecializing, pass infx.PH for values that shouldn’t bespecialized. For example:

deff(x):out=0forvinx.values():out+=vreturnoutf=fx.symbolic_trace(f,concrete_args={"x":{"a":fx.PH,"b":fx.PH,"c":fx.PH}})assertf({"a":1,"b":2,"c":4})==7
Parameters

  • root (Union[torch.nn.Module,Callable]) – Module or function to be traced and convertedinto a Graph representation.

  • concrete_args (Optional[Dict[str,any]]) – Inputs to be partially specialized

Returns

a Module created from the recorded operations fromroot.

Return type

GraphModule

Note

Backwards-compatibility for this API is guaranteed.

torch.fx.wrap(fn_or_name)[source]#

This function can be called at module-level scope to register fn_or_name as a “leaf function”.A “leaf function” will be preserved as a CallFunction node in the FX trace instead of beingtraced through:

# foo/bar/baz.pydefmy_custom_function(x,y):returnx*x+y*ytorch.fx.wrap("my_custom_function")deffn_to_be_traced(x,y):# When symbolic tracing, the below call to my_custom_function will be inserted into# the graph rather than tracing it.returnmy_custom_function(x,y)

This function can also equivalently be used as a decorator:

# foo/bar/baz.py@torch.fx.wrapdefmy_custom_function(x,y):returnx*x+y*y

A wrapped function can be thought of a “leaf function”, analogous to the concept of“leaf modules”, that is, they are functions that are left as calls in the FX tracerather than traced through.

Parameters

fn_or_name (Union[str,Callable]) – The function or name of the global function to insert into thegraph when it’s called

Note

Backwards-compatibility for this API is guaranteed.

classtorch.fx.GraphModule(*args,**kwargs)[source]#

GraphModule is an nn.Module generated from an fx.Graph. Graphmodule has agraph attribute, as well ascode andforward attributes generatedfrom thatgraph.

Warning

Whengraph is reassigned,code andforward will be automaticallyregenerated. However, if you edit the contents of thegraph without reassigningthegraph attribute itself, you must callrecompile() to update the generatedcode.

Note

Backwards-compatibility for this API is guaranteed.

__init__(root,graph,class_name='GraphModule')[source]#

Construct a GraphModule.

Parameters

  • root (Union[torch.nn.Module,Dict[str,Any]) –root can either be an nn.Module instance or a Dict mapping strings to any attribute type.In the case thatroot is a Module, any references to Module-based objects (via qualifiedname) in the Graph’s Nodes’target field will be copied over from the respective placewithinroot’s Module hierarchy into the GraphModule’s module hierarchy.In the case thatroot is a dict, the qualified name found in a Node’starget will belooked up directly in the dict’s keys. The object mapped to by the Dict will be copiedover into the appropriate place within the GraphModule’s module hierarchy.

  • graph (Graph) –graph contains the nodes this GraphModule should use for code generation

  • class_name (str) –name denotes the name of this GraphModule for debugging purposes. If it’s unset, allerror messages will report as originating fromGraphModule. It may be helpful to set thistoroot’s original name or a name that makes sense within the context of your transform.

Note

Backwards-compatibility for this API is guaranteed.

add_submodule(target,m)[source]#

Adds the given submodule toself.

This installs empty Modules where none exist yet if they aresubpaths oftarget.

Parameters

  • target (str) – The fully-qualified string name of the new submodule(See example innn.Module.get_submodule for how tospecify a fully-qualified string.)

  • m (Module) – The submodule itself; the actual object we want toinstall in the current Module

Returns

Whether or not the submodule could be inserted. For

this method to return True, each object in the chaindenoted bytarget must either a) not exist yet,or b) reference annn.Module (not a parameter orother attribute)

Return type

bool

Note

Backwards-compatibility for this API is guaranteed.

propertycode:str#

Return the Python code generated from theGraph underlying thisGraphModule.

delete_all_unused_submodules()[source]#

Deletes all unused submodules fromself.

A Module is considered “used” if any one of the following istrue:1. It has children that are used2. Its forward is called directly via acall_module node3. It has a non-Module attribute that is used from aget_attr node

This method can be called to clean up annn.Module withoutmanually callingdelete_submodule on each unused submodule.

Note

Backwards-compatibility for this API is guaranteed.

delete_submodule(target)[source]#

Deletes the given submodule fromself.

The module will not be deleted iftarget is not a validtarget.

Parameters

target (str) – The fully-qualified string name of the new submodule(See example innn.Module.get_submodule for how tospecify a fully-qualified string.)

Returns

Whether or not the target string referenced a

submodule we want to delete. A return value ofFalsemeans that thetarget was not a valid reference toa submodule.

Return type

bool

Note

Backwards-compatibility for this API is guaranteed.

propertygraph:Graph#

Return theGraph underlying thisGraphModule

print_readable(print_output=True,include_stride=False,include_device=False,colored=False,*,fast_sympy_print=False,expanded_def=False)[source]#

Return the Python code generated for current GraphModule and its children GraphModules

Warning

This API is experimental and isNOT backward-compatible.

recompile()[source]#

Recompile this GraphModule from itsgraph attribute. This should becalled after editing the containedgraph, otherwise the generatedcode of thisGraphModule will be out of date.

Note

Backwards-compatibility for this API is guaranteed.

Return type

PythonCode

to_folder(folder,module_name='FxModule')[source]#
Dumps out module tofolder withmodule_name so that it can be

imported withfrom<folder>import<module_name>

Args:

folder (Union[str, os.PathLike]): The folder to write the code out to

module_name (str): Top-level name to use for theModule while

writing out the code

Warning

This API is experimental and isNOT backward-compatible.

classtorch.fx.Graph(owning_module=None,tracer_cls=None,tracer_extras=None)[source]#

Graph is the main data structure used in the FX Intermediate Representation.It consists of a series ofNode s, each representing callsites (or othersyntactic constructs). The list ofNode s, taken together, constitute avalid Python function.

For example, the following code

importtorchimporttorch.fxclassMyModule(torch.nn.Module):def__init__(self):super().__init__()self.param=torch.nn.Parameter(torch.rand(3,4))self.linear=torch.nn.Linear(4,5)defforward(self,x):returntorch.topk(torch.sum(self.linear(x+self.linear.weight).relu(),dim=-1),3)m=MyModule()gm=torch.fx.symbolic_trace(m)

Will produce the following Graph:

print(gm.graph)
graph(x):    %linear_weight : [num_users=1] = self.linear.weight    %add_1 : [num_users=1] = call_function[target=operator.add](args = (%x, %linear_weight), kwargs = {})    %linear_1 : [num_users=1] = call_module[target=linear](args = (%add_1,), kwargs = {})    %relu_1 : [num_users=1] = call_method[target=relu](args = (%linear_1,), kwargs = {})    %sum_1 : [num_users=1] = call_function[target=torch.sum](args = (%relu_1,), kwargs = {dim: -1})    %topk_1 : [num_users=1] = call_function[target=torch.topk](args = (%sum_1, 3), kwargs = {})    return topk_1

For the semantics of operations represented in theGraph, please seeNode.

Note

Backwards-compatibility for this API is guaranteed.

__init__(owning_module=None,tracer_cls=None,tracer_extras=None)[source]#

Construct an empty Graph.

Note

Backwards-compatibility for this API is guaranteed.

call_function(the_function,args=None,kwargs=None,type_expr=None,name=None)[source]#

Insert acall_functionNode into theGraph. Acall_function noderepresents a call to a Python callable, specified bythe_function.

Parameters

  • the_function (Callable[...,Any]) – The function to be called. Can be any PyTorchoperator, Python function, or member of thebuiltins oroperatornamespaces.

  • args (Optional[Tuple[Argument,...]]) – The positional arguments to be passedto the called function.

  • kwargs (Optional[Dict[str,Argument]]) – The keyword arguments to be passedto the called function

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

  • name (Optional[str]) – The name of the node. If not specified, set to None

Returns

The newly created and insertedcall_function node.

Return type

Node

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node().

Note

Backwards-compatibility for this API is guaranteed.

call_method(method_name,args=None,kwargs=None,type_expr=None)[source]#

Insert acall_methodNode into theGraph. Acall_method noderepresents a call to a given method on the 0th element ofargs.

Parameters

  • method_name (str) – The name of the method to apply to the self argument.For example, if args[0] is aNode representing aTensor,then to callrelu() on thatTensor, passrelu tomethod_name.

  • args (Optional[Tuple[Argument,...]]) – The positional arguments to be passedto the called method. Note that thisshould include aself argument.

  • kwargs (Optional[Dict[str,Argument]]) – The keyword arguments to be passedto the called method

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

Returns

The newly created and insertedcall_method node.

Return type

Node

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node().

Note

Backwards-compatibility for this API is guaranteed.

call_module(module_name,args=None,kwargs=None,type_expr=None)[source]#

Insert acall_moduleNode into theGraph. Acall_module noderepresents a call to the forward() function of aModule in theModulehierarchy.

Parameters

  • module_name (str) – The qualified name of theModule in theModulehierarchy to be called. For example, if the tracedModule has asubmodule namedfoo, which has a submodule namedbar, thequalified namefoo.bar should be passed asmodule_name tocall that module.

  • args (Optional[Tuple[Argument,...]]) – The positional arguments to be passedto the called method. Note that this shouldnot include aself argument.

  • kwargs (Optional[Dict[str,Argument]]) – The keyword arguments to be passedto the called method

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

Returns

The newly-created and insertedcall_module node.

Return type

Node

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node().

Note

Backwards-compatibility for this API is guaranteed.

create_node(op,target,args=None,kwargs=None,name=None,type_expr=None)[source]#

Create aNode and add it to theGraph at the current insert-point.Note that the current insert-point can be set viaGraph.inserting_before()andGraph.inserting_after().

Parameters

  • op (str) – the opcode for this Node. One of ‘call_function’, ‘call_method’, ‘get_attr’,‘call_module’, ‘placeholder’, or ‘output’. The semantics of these opcodes aredescribed in theGraph docstring.

  • args (Optional[Tuple[Argument,...]]) – is a tuple of arguments to this node.

  • kwargs (Optional[Dict[str,Argument]]) – the kwargs of this Node

  • name (Optional[str]) – an optional string name for theNode.This will influence the name of the value assigned to in thePython generated code.

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

Returns

The newly-created and inserted node.

Return type

Node

Note

Backwards-compatibility for this API is guaranteed.

eliminate_dead_code(is_impure_node=None)[source]#

Remove all dead code from the graph, based on each node’s number ofusers, and whether the nodes have any side effects. The graph must betopologically sorted before calling.

Parameters

  • is_impure_node (Optional[Callable[[Node],bool]]) – A function that returns

  • None (whether a node is impure. If this is) –

  • to (then the default behavior is) –

  • Node.is_impure. (use) –

Returns

Whether the graph was changed as a result of the pass.

Return type

bool

Example:

Before dead code is eliminated,a froma = x + 1 below has no usersand thus can be eliminated from the graph without having an effect.

defforward(self,x):a=x+1returnx+self.attr_1

After dead code is eliminated,a = x + 1 has been removed, and the restofforward remains.

defforward(self,x):returnx+self.attr_1

Warning

Dead code elimination has some heuristics to avoid removingside-effectful nodes (see Node.is_impure) but in general coverageis very bad, so you should assume that this method is not soundto call unless you know that your FX graph consists entirelyof functional operations or you supply your own customfunction for detecting side-effectful nodes.

Note

Backwards-compatibility for this API is guaranteed.

erase_node(to_erase)[source]#

Erases aNode from theGraph. Throws an exception ifthere are still users of that node in theGraph.

Parameters

to_erase (Node) – TheNode to erase from theGraph.

Note

Backwards-compatibility for this API is guaranteed.

find_nodes(*,op,target=None,sort=True)[source]#

Allows for fast query of nodes

Parameters

  • op (str) – the name of the operation

  • target (Optional[Target]) – the target of the node. For call_function,the target is required. For other ops, the target is optional.

  • sort (bool) – whether to return nodes in the order they appear onon the graph.

Returns

Iterable of nodes with the requested op and target.

Warning

This API is experimental and isNOT backward-compatible.

get_attr(qualified_name,type_expr=None)[source]#

Insert aget_attr node into the Graph. Aget_attrNode represents thefetch of an attribute from theModule hierarchy.

Parameters

  • qualified_name (str) – the fully-qualified name of the attribute to be retrieved.For example, if the traced Module has a submodule namedfoo, which has asubmodule namedbar, which has an attribute namedbaz, the qualifiednamefoo.bar.baz should be passed asqualified_name.

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

Returns

The newly-created and insertedget_attr node.

Return type

Node

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node.

Note

Backwards-compatibility for this API is guaranteed.

graph_copy(g,val_map,return_output_node=False)[source]#

Copy all nodes from a given graph intoself.

Parameters

  • g (Graph) – The source graph from which to copy Nodes.

  • val_map (Dict[Node,Node]) – a dictionary that will be populated with a mappingfrom nodes ing to nodes inself. Note thatval_map can be passedin with values in it already to override copying of certain values.

Returns

The value inself that is now equivalent to the output value ing,ifg had anoutput node.None otherwise.

Return type

Optional[Union[tuple[‘Argument’, …],Sequence[Argument],Mapping[str, Argument],slice,range,Node,str,int,float,bool,complex,dtype,Tensor,device,memory_format,layout,OpOverload,SymInt,SymBool,SymFloat]]

Note

Backwards-compatibility for this API is guaranteed.

inserting_after(n=None)[source]#
Set the point at which create_node and companion methods will insert into the graph.

When used within a ‘with’ statement, this will temporary set the insert point andthen restore it when the with statement exits:

withg.inserting_after(n):...# inserting after node n...# insert point restored to what it was previouslyg.inserting_after(n)#  set the insert point permanently

Args:

n (Optional[Node]): The node before which to insert. If None this will insert after

the beginning of the entire graph.

Returns:

A resource manager that will restore the insert point on__exit__.

Note

Backwards-compatibility for this API is guaranteed.

inserting_before(n=None)[source]#
Set the point at which create_node and companion methods will insert into the graph.

When used within a ‘with’ statement, this will temporary set the insert point andthen restore it when the with statement exits:

withg.inserting_before(n):...# inserting before node n...# insert point restored to what it was previouslyg.inserting_before(n)#  set the insert point permanently

Args:

n (Optional[Node]): The node before which to insert. If None this will insert before

the beginning of the entire graph.

Returns:

A resource manager that will restore the insert point on__exit__.

Note

Backwards-compatibility for this API is guaranteed.

lint()[source]#

Runs various checks on this Graph to make sure it is well-formed. Inparticular:- Checks Nodes have correct ownership (owned by this graph)- Checks Nodes appear in topological order- If this Graph has an owning GraphModule, checks that targetsexist in that GraphModule

Note

Backwards-compatibility for this API is guaranteed.

node_copy(node,arg_transform=<functionGraph.<lambda>>)[source]#

Copy a node from one graph into another.arg_transform needs to transform arguments fromthe graph of node to the graph of self. Example:

# Copying all the nodes in `g` into `new_graph`g:torch.fx.Graph=...new_graph=torch.fx.graph()value_remap={}fornodeing.nodes:value_remap[node]=new_graph.node_copy(node,lambdan:value_remap[n])
Parameters

  • node (Node) – The node to copy intoself.

  • arg_transform (Callable[[Node],Argument]) – A function that transformsNode arguments in node’sargs andkwargs into theequivalent argument inself. In the simplest case, this shouldretrieve a value out of a table mapping Nodes in the originalgraph toself.

Return type

Node

Note

Backwards-compatibility for this API is guaranteed.

propertynodes:_node_list#

Get the list of Nodes that constitute this Graph.

Note that thisNode list representation is a doubly-linked list. Mutationsduring iteration (e.g. delete a Node, add a Node) are safe.

Returns

A doubly-linked list of Nodes. Note thatreversed can be called onthis list to switch iteration order.

on_generate_code(make_transformer)[source]#

Register a transformer function when python code is generated

Args:
make_transformer (Callable[[Optional[TransformCodeFunc]], TransformCodeFunc]):

a function that returns a code transformer to be registered.This function is called byon_generate_code to obtain thecode transformer.

This function is also given as its input the currentlyregistered code transformer (or None if nothing is registered),in case it is not desirable to overwrite it. This is useful tochain code transformers together.

Returns:

a context manager that when used in awith statement, to automaticallyrestore the previously registered code transformer.

Example:

gm:fx.GraphModule=...# This is a code transformer we want to register. This code# transformer prepends a pdb import and trace statement at the very# beginning of the generated torch.fx code to allow for manual# debugging with the PDB library.definsert_pdb(body):return["import pdb; pdb.set_trace()\n",*body]# Registers `insert_pdb`, and overwrites the current registered# code transformer (given by `_` to the lambda):gm.graph.on_generate_code(lambda_:insert_pdb)# Or alternatively, registers a code transformer which first# runs `body` through existing registered transformer, then# through `insert_pdb`:gm.graph.on_generate_code(lambdacurrent_trans:(lambdabody:insert_pdb(current_trans(body)ifcurrent_transelsebody)))gm.recompile()gm(*inputs)# drops into pdb

This function can also be used as a context manager, with the benefit toautomatically restores the previously registered code transformer:

# ... continue from previous examplewithgm.graph.on_generate_code(lambda_:insert_pdb):# do more stuff with `gm`...gm.recompile()gm(*inputs)# drops into pdb# now previous code transformer is restored (but `gm`'s code with pdb# remains - that means you can run `gm` with pdb here too, until you# run next `recompile()`).

Warning

This API is experimental and isNOT backward-compatible.

output(result,type_expr=None)[source]#

Insert anoutputNode into theGraph. Anoutput node representsareturn statement in Python code.result is the value that shouldbe returned.

Parameters

  • result (Argument) – The value to be returned.

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have.

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node.

Note

Backwards-compatibility for this API is guaranteed.

output_node()[source]#

Warning

This API is experimental and isNOT backward-compatible.

Return type

Node

placeholder(name,type_expr=None,default_value)[source]#

Insert aplaceholder node into the Graph. Aplaceholder representsa function input.

Parameters

  • name (str) – A name for the input value. This corresponds to the nameof the positional argument to the function thisGraph represents.

  • type_expr (Optional[Any]) – an optional type annotation representing thePython type the output of this node will have. This is needed in somecases for proper code generation (e.g. when the function is usedsubsequently in TorchScript compilation).

  • default_value (Any) – The default value this function argument should takeon. NOTE: to allow forNone as a default value,inspect.Signature.emptyshould be passed as this argument to specify that the parameter does _not_have a default value.

Return type

Node

Note

The same insertion point and type expression rules apply for this methodasGraph.create_node.

Note

Backwards-compatibility for this API is guaranteed.

print_tabular()[source]#

Prints the intermediate representation of the graph in tabularformat. Note that this API requires thetabulate module to beinstalled.

Note

Backwards-compatibility for this API is guaranteed.

process_inputs(*args)[source]#

Processes args so that they can be passed to the FX graph.

Warning

This API is experimental and isNOT backward-compatible.

process_outputs(out)[source]#

Warning

This API is experimental and isNOT backward-compatible.

python_code(root_module,*,verbose=False,include_stride=False,include_device=False,colored=False,expanded_def=False)[source]#

Turn thisGraph into valid Python code.

Parameters

root_module (str) – The name of the root module on which to look-upqualified name targets. This is usually ‘self’.

Returns

src: the Python source code representing the objectglobals: a dictionary of global names insrc -> the objects that they reference.

Return type

A PythonCode object, consisting of two fields

Note

Backwards-compatibility for this API is guaranteed.

set_codegen(codegen)[source]#

Warning

This API is experimental and isNOT backward-compatible.

classtorch.fx.Node(graph,name,op,target,args,kwargs,return_type=None)[source]#

Node is the data structure that represents individual operations withinaGraph. For the most part, Nodes represent callsites to various entities,such as operators, methods, and Modules (some exceptions include nodes thatspecify function inputs and outputs). EachNode has a function specifiedby itsop property. TheNode semantics for each value ofop are as follows:

  • placeholder represents a function input. Thename attribute specifies the name this value will take on.target is similarly the name of the argument.args holds either: 1) nothing, or 2) a single argumentdenoting the default parameter of the function input.kwargs is don’t-care. Placeholders correspond tothe function parameters (e.g.x) in the graph printout.

  • get_attr retrieves a parameter from the module hierarchy.name is similarly the name the result of thefetch is assigned to.target is the fully-qualified name of the parameter’s position in the module hierarchy.args andkwargs are don’t-care

  • call_function applies a free function to some values.name is similarly the name of the value to assignto.target is the function to be applied.args andkwargs represent the arguments to the function,following the Python calling convention

  • call_module applies a module in the module hierarchy’sforward() method to given arguments.name isas previous.target is the fully-qualified name of the module in the module hierarchy to call.args andkwargs represent the arguments to invoke the module on,excluding the self argument.

  • call_method calls a method on a value.name is as similar.target is the string name of the methodto apply to theself argument.args andkwargs represent the arguments to invoke the module on,including the self argument

  • output contains the output of the traced function in itsargs[0] attribute. This corresponds to the “return” statementin the Graph printout.

Note

Backwards-compatibility for this API is guaranteed.

propertyall_input_nodes:list['Node']#

Return all Nodes that are inputs to this Node. This is equivalent toiterating overargs andkwargs and only collecting the values thatare Nodes.

Returns

List ofNodes that appear in theargs andkwargs of thisNode, in that order.

append(x)[source]#

Insertx after this node in the list of nodes in the graph.Equivalent toself.next.prepend(x)

Parameters

x (Node) – The node to put after this node. Must be a member of the same graph.

Note

Backwards-compatibility for this API is guaranteed.

propertyargs:tuple[Union[tuple['Argument',...],collections.abc.Sequence['Argument'],collections.abc.Mapping[str,'Argument'],slice,range,torch.fx.node.Node,str,int,float,bool,complex,torch.dtype,torch.Tensor,torch.device,torch.memory_format,torch.layout,torch._ops.OpOverload,torch.SymInt,torch.SymBool,torch.SymFloat,NoneType],...]#

The tuple of arguments to thisNode. The interpretation of argumentsdepends on the node’s opcode. See theNode docstring for moreinformation.

Assignment to this property is allowed. All accounting of uses and usersis updated automatically on assignment.

format_node(placeholder_names=None,maybe_return_typename=None,*,include_tensor_metadata=False)[source]#

Return a descriptive string representation ofself.

This method can be used with no arguments as a debuggingutility.

This function is also used internally in the__str__ methodofGraph. Together, the strings inplaceholder_namesandmaybe_return_typename make up the signature of theautogeneratedforward function in this Graph’s surroundingGraphModule.placeholder_names andmaybe_return_typenameshould not be used otherwise.

Parameters

  • placeholder_names (Optional[list[str]]) – A list that will store formatted stringsrepresenting the placeholders in the generatedforward function. Internal use only.

  • maybe_return_typename (Optional[list[str]]) – A single-element list that will storea formatted string representing the output of thegeneratedforward function. Internal use only.

  • include_tensor_metadata (bool) – Whether to include tensor metadata

Returns

If 1) we’re usingformat_node as an internal helper

in the__str__ method ofGraph, and 2)selfis a placeholder Node, returnNone. Otherwise,return a descriptive string representation of thecurrent Node.

Return type

str

Note

Backwards-compatibility for this API is guaranteed.

insert_arg(idx,arg)[source]#

Insert an positional argument to the argument list with given index.

Parameters

  • idx (int) – The index of the element inself.args to be inserted before.

  • arg (Argument) – The new argument value to insert intoargs

Note

Backwards-compatibility for this API is guaranteed.

is_impure(impure_random=True)[source]#

Returns whether this op is impure, i.e. if its op is a placeholder oroutput, or if a call_function or call_module which is impure.

Parameters

impure_random (bool) – Whether to treat rand op as impure.

Returns

If the op is impure or not.

Return type

bool

Warning

This API is experimental and isNOT backward-compatible.

propertykwargs:dict[str,Union[tuple['Argument',...],collections.abc.Sequence['Argument'],collections.abc.Mapping[str,'Argument'],slice,range,torch.fx.node.Node,str,int,float,bool,complex,torch.dtype,torch.Tensor,torch.device,torch.memory_format,torch.layout,torch._ops.OpOverload,torch.SymInt,torch.SymBool,torch.SymFloat,NoneType]]#

The dict of keyword arguments to thisNode. The interpretation of argumentsdepends on the node’s opcode. See theNode docstring for moreinformation.

Assignment to this property is allowed. All accounting of uses and usersis updated automatically on assignment.

propertynext:Node#

Returns the nextNode in the linked list of Nodes.

Returns

The nextNode in the linked list of Nodes.

normalized_arguments(root,arg_types=None,kwarg_types=None,normalize_to_only_use_kwargs=False)[source]#

Returns normalized arguments to Python targets. This means thatargs/kwargs will be matched up to the module/functional’ssignature and return exclusively kwargs in positional orderifnormalize_to_only_use_kwargs is true.Also populates default values. Does not support positional-onlyparameters or varargs parameters.

Supports module calls.

May requirearg_types andkwarg_types in order to disambiguate overloads.

Parameters

  • root (torch.nn.Module) – Module upon which to resolve module targets.

  • arg_types (Optional[Tuple[Any]]) – Tuple of arg types for the args

  • kwarg_types (Optional[Dict[str,Any]]) – Dict of arg types for the kwargs

  • normalize_to_only_use_kwargs (bool) – Whether to normalize to only use kwargs.

Returns

Returns NamedTuple ArgsKwargsPair, orNone if not successful.

Return type

Optional[ArgsKwargsPair]

Warning

This API is experimental and isNOT backward-compatible.

prepend(x)[source]#

Insert x before this node in the list of nodes in the graph. Example:

Before:p->selfbx->x->axAfter:p->x->selfbx->ax
Parameters

x (Node) – The node to put before this node. Must be a member of the same graph.

Note

Backwards-compatibility for this API is guaranteed.

propertyprev:Node#

Returns the previousNode in the linked list of Nodes.

Returns

The previousNode in the linked list of Nodes.

replace_all_uses_with(replace_with,delete_user_cb=<functionNode.<lambda>>,*,propagate_meta=False)[source]#

Replace all uses ofself in the Graph with the Nodereplace_with.

Parameters

  • replace_with (Node) – The node to replace all uses ofself with.

  • delete_user_cb (Callable) – Callback that is called to determinewhether a given user of the self node should be removed.

  • propagate_meta (bool) – Whether or not to copy all propertieson the .meta field of the original node onto the replacement node.For safety, this is only valid to do if the replacement nodedoesn’t already have an existing .meta field.

Returns

The list of Nodes on which this change was made.

Return type

list[‘Node’]

Note

Backwards-compatibility for this API is guaranteed.

replace_input_with(old_input,new_input)[source]#

Loop through input nodes ofself, and replace all instances ofold_input withnew_input.

Parameters

  • old_input (Node) – The old input node to be replaced.

  • new_input (Node) – The new input node to replaceold_input.

Note

Backwards-compatibility for this API is guaranteed.

propertystack_trace:Optional[str]#

Return the Python stack trace that was recorded during tracing, if any.When traced with fx.Tracer, this property is usually populated byTracer.create_proxy. To record stack traces during tracing for debug purposes,setrecord_stack_traces = True on theTracer instance.When traced with dynamo, this property will be populated by default byOutputGraph.create_proxy.

stack_trace would have the innermost frame at the end of the string.

update_arg(idx,arg)[source]#

Update an existing positional argument to contain the new valuearg. After calling,self.args[idx]==arg.

Parameters

  • idx (int) – The index intoself.args of the element to update

  • arg (Argument) – The new argument value to write intoargs

Note

Backwards-compatibility for this API is guaranteed.

update_kwarg(key,arg)[source]#

Update an existing keyword argument to contain the new valuearg. After calling,self.kwargs[key]==arg.

Parameters

  • key (str) – The key inself.kwargs of the element to update

  • arg (Argument) – The new argument value to write intokwargs

Note

Backwards-compatibility for this API is guaranteed.

classtorch.fx.Tracer(autowrap_modules=(math,),autowrap_functions=())[source]#

Tracer is the class that implements the symbolic tracing functionalityoftorch.fx.symbolic_trace. A call tosymbolic_trace(m) is equivalenttoTracer().trace(m).

Tracer can be subclassed to override various behaviors of the tracingprocess. The different behaviors that can be overridden are describedin the docstrings of the methods on this class.

Note

Backwards-compatibility for this API is guaranteed.

call_module(m,forward,args,kwargs)[source]#

Method that specifies the behavior of thisTracer when it encountersa call to annn.Module instance.

By default, the behavior is to check if the called module is a leaf moduleviais_leaf_module. If it is, emit acall_module node referring tom in theGraph. Otherwise, call theModule normally, tracing throughthe operations in itsforward function.

This method can be overridden to–for example–create nested tracedGraphModules, or any other behavior you would want while tracing acrossModule boundaries.

Parameters

  • m (Module) – The module for which a call is being emitted

  • forward (Callable) – The forward() method of theModule to be invoked

  • args (Tuple) – args of the module callsite

  • kwargs (Dict) – kwargs of the module callsite

Returns

The return value from the Module call. In the case that acall_modulenode was emitted, this is aProxy value. Otherwise, it is whatevervalue was returned from theModule invocation.

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

create_arg(a)[source]#

A method to specify the behavior of tracing when preparing values tobe used as arguments to nodes in theGraph.

By default, the behavior includes:

  1. Iterate through collection types (e.g. tuple, list, dict) and recursivelycallcreate_args on the elements.

  2. Given a Proxy object, return a reference to the underlying IRNode

  3. Given a non-Proxy Tensor object, emit IR for various cases:

    • For a Parameter, emit aget_attr node referring to that Parameter

    • For a non-Parameter Tensor, store the Tensor away in a specialattribute referring to that attribute.

This method can be overridden to support more types.

Parameters

a (Any) – The value to be emitted as anArgument in theGraph.

Returns

The valuea converted into the appropriateArgument

Return type

Argument

Note

Backwards-compatibility for this API is guaranteed.

create_args_for_root(root_fn,is_module,concrete_args=None)[source]#

Createplaceholder nodes corresponding to the signature of therootModule. This method introspects root’s signature and emits thosenodes accordingly, also supporting*args and**kwargs.

Warning

This API is experimental and isNOT backward-compatible.

create_node(kind,target,args,kwargs,name=None,type_expr=None)[source]#

Inserts a graph node given target, args, kwargs, and name.

This method can be overridden to do extra checking, validation, ormodification of values used in node creation. For example, one mightwant to disallow in-place operations from being recorded.

Note

Backwards-compatibility for this API is guaranteed.

Return type

Node

create_proxy(kind,target,args,kwargs,name=None,type_expr=None,proxy_factory_fn=None)[source]#

Create a Node from the given arguments, then return the Nodewrapped in a Proxy object.

If kind = ‘placeholder’, then we’re creating a Node thatrepresents the parameter of a function. If we need to encodea default parameter, we use theargs tuple.args isotherwise empty forplaceholder Nodes.

Note

Backwards-compatibility for this API is guaranteed.

get_fresh_qualname(prefix)[source]#

Gets a fresh name for a prefix and returns it. This function ensuresthat it will not clash with an existing attribute on the graph.

Note

Backwards-compatibility for this API is guaranteed.

Return type

str

getattr(attr,attr_val,parameter_proxy_cache)[source]#

Method that specifies the behavior of thisTracer when we call getattron a call to annn.Module instance.

By default, the behavior is to return a proxy value for the attribute. Italso stores the proxy value in theparameter_proxy_cache, so that futurecalls will reuse the proxy rather than creating a new one.

This method can be overridden to –for example– not return proxies whenquerying parameters.

Parameters

  • attr (str) – The name of the attribute being queried

  • attr_val (Any) – The value of the attribute

  • parameter_proxy_cache (Dict[str,Any]) – A cache of attr names to proxies

Returns

The return value from the getattr call.

Warning

This API is experimental and isNOT backward-compatible.

is_leaf_module(m,module_qualified_name)[source]#

A method to specify whether a givennn.Module is a “leaf” module.

Leaf modules are the atomic units that appear inthe IR, referenced bycall_module calls. By default,Modules in the PyTorch standard library namespace (torch.nn)are leaf modules. All other modules are traced through andtheir constituent ops are recorded, unless specified otherwisevia this parameter.

Parameters

  • m (Module) – The module being queried about

  • module_qualified_name (str) – The path to root of this module. For example,if you have a module hierarchy where submodulefoo containssubmodulebar, which contains submodulebaz, that module willappear with the qualified namefoo.bar.baz here.

Return type

bool

Note

Backwards-compatibility for this API is guaranteed.

iter(obj)[source]#
Called when a proxy object is being iterated over, such as

when used in control flow. Normally we don’t know what to do becausewe don’t know the value of the proxy, but a custom tracer can attach moreinformation to the graph node using create_node and can choose to return an iterator.

Note

Backwards-compatibility for this API is guaranteed.

Return type

Iterator

keys(obj)[source]#
Called when a proxy object is has the keys() method called.

This is what happens when ** is called on a proxy. This should return aniterator it ** is suppose to work in your custom tracer.

Note

Backwards-compatibility for this API is guaranteed.

Return type

Any

path_of_module(mod)[source]#

Helper method to find the qualified name ofmod in the Module hierarchyofroot. For example, ifroot has a submodule namedfoo, which hasa submodule namedbar, passingbar into this function will returnthe string “foo.bar”.

Parameters

mod (str) – TheModule to retrieve the qualified name for.

Return type

str

Note

Backwards-compatibility for this API is guaranteed.

proxy(node)[source]#

Note

Backwards-compatibility for this API is guaranteed.

Return type

Proxy

to_bool(obj)[source]#
Called when a proxy object is being converted to a boolean, such as

when used in control flow. Normally we don’t know what to do becausewe don’t know the value of the proxy, but a custom tracer can attach moreinformation to the graph node using create_node and can choose to return a value.

Note

Backwards-compatibility for this API is guaranteed.

Return type

bool

trace(root,concrete_args=None)[source]#

Traceroot and return the corresponding FXGraph representation.rootcan either be annn.Module instance or a Python callable.

Note that after this call,self.root may be different from theroot passedin here. For example, when a free function is passed totrace(), we willcreate annn.Module instance to use as the root and add embedded constantsto.

Parameters

  • root (Union[Module,Callable]) – Either aModule or a function to betraced through. Backwards-compatibility for this parameter isguaranteed.

  • concrete_args (Optional[Dict[str,any]]) – Concrete arguments that shouldnot be treated as Proxies. This parameter is experimental andits backwards-compatibility isNOT guaranteed.

Returns

AGraph representing the semantics of the passed-inroot.

Return type

Graph

Note

Backwards-compatibility for this API is guaranteed.

classtorch.fx.Proxy(node,tracer=None)[source]#

Proxy objects areNode wrappers that flow through theprogram during symbolic tracing and record all the operations(torch function calls, method calls, operators) that they touchinto the growing FX Graph.

If you’re doing graph transforms, you can wrap your ownProxymethod around a rawNode so that you can use the overloadedoperators to add additional things to aGraph.

Proxy objects cannot be iterated. In other words, the symbolictracer will throw an error if aProxy is used in a loop or asan*args/**kwargs function argument.

There are two main ways around this:1. Factor out the untraceable logic into a top-level function andusefx.wrap on it.2. If the control flow is static (i.e. the loop trip count isbased on some hyperparameter), the code can be kept in its originalposition and refactored into something like:

foriinrange(self.some_hyperparameter):indexed_item=proxied_value[i]

For a more detailed description into the Proxy internals, check outthe “Proxy” section intorch/fx/README.md

Note

Backwards-compatibility for this API is guaranteed.

classtorch.fx.Interpreter(module,garbage_collect_values=True,graph=None)[source]#

An Interpreter executes an FX graph Node-by-Node. This patterncan be useful for many things, including writing codetransformations as well as analysis passes.

Methods in the Interpreter class can be overridden to customizethe behavior of execution. The map of overridable methodsin terms of call hierarchy:

run()+--run_node+--placeholder()+--get_attr()+--call_function()+--call_method()+--call_module()+--output()

Example

Suppose we want to swap all instances oftorch.neg withtorch.sigmoid and vice versa (including theirTensormethod equivalents). We could subclass Interpreter like so:

classNegSigmSwapInterpreter(Interpreter):defcall_function(self,target:Target,args:Tuple,kwargs:Dict)->Any:iftarget==torch.sigmoid:returntorch.neg(*args,**kwargs)returnsuper().call_function(target,args,kwargs)defcall_method(self,target:Target,args:Tuple,kwargs:Dict)->Any:iftarget=="neg":call_self,*args_tail=argsreturncall_self.sigmoid(*args_tail,**kwargs)returnsuper().call_method(target,args,kwargs)deffn(x):returntorch.sigmoid(x).neg()gm=torch.fx.symbolic_trace(fn)input=torch.randn(3,4)result=NegSigmSwapInterpreter(gm).run(input)torch.testing.assert_close(result,torch.neg(input).sigmoid())
Parameters

  • module (torch.nn.Module) – The module to be executed

  • garbage_collect_values (bool) – Whether to delete values after their lastuse within the Module’s execution. This ensures optimal memory usage duringexecution. This can be disabled to, for example, examine all of the intermediatevalues in the execution by looking at theInterpreter.env attribute.

  • graph (Optional[Graph]) – If passed, the interpreter will execute thisgraph instead ofmodule.graph, using the providedmoduleargument to satisfy any requests for state.

Note

Backwards-compatibility for this API is guaranteed.

boxed_run(args_list)[source]#

Runmodule via interpretation and return the result. This uses the “boxed”calling convention, where you pass a list of arguments, which will be clearedby the interpreter. This ensures that input tensors are promptly deallocated.

Note

Backwards-compatibility for this API is guaranteed.

call_function(target,args,kwargs)[source]#

Execute acall_function node and return the result.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Return type

Any

Return

Any: The value returned by the function invocation

Note

Backwards-compatibility for this API is guaranteed.

call_method(target,args,kwargs)[source]#

Execute acall_method node and return the result.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Return type

Any

Return

Any: The value returned by the method invocation

Note

Backwards-compatibility for this API is guaranteed.

call_module(target,args,kwargs)[source]#

Execute acall_module node and return the result.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Return type

Any

Return

Any: The value returned by the module invocation

Note

Backwards-compatibility for this API is guaranteed.

fetch_args_kwargs_from_env(n)[source]#

Fetch the concrete values ofargs andkwargs of nodenfrom the current execution environment.

Parameters

n (Node) – The node for whichargs andkwargs should be fetched.

Returns

args andkwargs with concrete values forn.

Return type

Tuple[Tuple, Dict]

Note

Backwards-compatibility for this API is guaranteed.

fetch_attr(target)[source]#

Fetch an attribute from theModule hierarchy ofself.module.

Parameters

target (str) – The fully-qualified name of the attribute to fetch

Returns

The value of the attribute.

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

get_attr(target,args,kwargs)[source]#

Execute aget_attr node. Will retrieve an attributevalue from theModule hierarchy ofself.module.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Returns

The value of the attribute that was retrieved

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

map_nodes_to_values(args,n)[source]#

Recursively descend throughargs and look up the concrete valuefor eachNode in the current execution environment.

Parameters

  • args (Argument) – Data structure within which to look up concrete values

  • n (Node) – Node to whichargs belongs. This is only used for error reporting.

Return type

Optional[Union[tuple[‘Argument’, …],Sequence[Argument],Mapping[str, Argument],slice,range,Node,str,int,float,bool,complex,dtype,Tensor,device,memory_format,layout,OpOverload,SymInt,SymBool,SymFloat]]

Note

Backwards-compatibility for this API is guaranteed.

output(target,args,kwargs)[source]#

Execute anoutput node. This really just retrievesthe value referenced by theoutput node and returns it.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Returns

The return value referenced by the output node

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

placeholder(target,args,kwargs)[source]#

Execute aplaceholder node. Note that this is stateful:Interpreter maintains an internal iterator overarguments passed torun and this method returnsnext() on that iterator.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Returns

The argument value that was retrieved.

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

run(*args,initial_env=None,enable_io_processing=True)[source]#

Runmodule via interpretation and return the result.

Parameters

  • *args – The arguments to the Module to run, in positional order

  • initial_env (Optional[Dict[Node,Any]]) – An optional starting environment for execution.This is a dict mappingNode to any value. This can be used, for example, topre-populate results for certainNodes so as to do only partial evaluation withinthe interpreter.

  • enable_io_processing (bool) – If true, we process the inputs and outputs with graph’s process_inputs andprocess_outputs function first before using them.

Returns

The value returned from executing the Module

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

run_node(n)[source]#

Run a specific noden and return the result.Calls into placeholder, get_attr, call_function,call_method, call_module, or output dependingonnode.op

Parameters

n (Node) – The Node to execute

Returns

The result of executingn

Return type

Any

Note

Backwards-compatibility for this API is guaranteed.

classtorch.fx.Transformer(module)[source]#

Transformer is a special type of interpreter that produces anewModule. It exposes atransform() method that returnsthe transformedModule.Transformer does not requirearguments to run, asInterpreter does.Transformer worksentirely symbolically.

Example

Suppose we want to swap all instances oftorch.neg withtorch.sigmoid and vice versa (including theirTensormethod equivalents). We could subclassTransformer like so:

classNegSigmSwapXformer(Transformer):defcall_function(self,target:"Target",args:Tuple[Argument,...],kwargs:Dict[str,Any],)->Any:iftarget==torch.sigmoid:returntorch.neg(*args,**kwargs)returnsuper().call_function(target,args,kwargs)defcall_method(self,target:"Target",args:Tuple[Argument,...],kwargs:Dict[str,Any],)->Any:iftarget=="neg":call_self,*args_tail=argsreturncall_self.sigmoid(*args_tail,**kwargs)returnsuper().call_method(target,args,kwargs)deffn(x):returntorch.sigmoid(x).neg()gm=torch.fx.symbolic_trace(fn)transformed:torch.nn.Module=NegSigmSwapXformer(gm).transform()input=torch.randn(3,4)torch.testing.assert_close(transformed(input),torch.neg(input).sigmoid())
Parameters

module (GraphModule) – TheModule to be transformed.

Note

Backwards-compatibility for this API is guaranteed.

call_function(target,args,kwargs)[source]#

Note

Backwards-compatibility for this API is guaranteed.

Return type

Any

call_module(target,args,kwargs)[source]#

Note

Backwards-compatibility for this API is guaranteed.

Return type

Any

get_attr(target,args,kwargs)[source]#

Execute aget_attr node. InTransformer, this isoverridden to insert a newget_attr node into the outputgraph.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Return type

Proxy

Note

Backwards-compatibility for this API is guaranteed.

placeholder(target,args,kwargs)[source]#

Execute aplaceholder node. InTransformer, this isoverridden to insert a newplaceholder into the outputgraph.

Parameters

  • target (Target) – The call target for this node. SeeNode fordetails on semantics

  • args (Tuple) – Tuple of positional args for this invocation

  • kwargs (Dict) – Dict of keyword arguments for this invocation

Return type

Proxy

Note

Backwards-compatibility for this API is guaranteed.

transform()[source]#

Transformself.module and return the transformedGraphModule.

Note

Backwards-compatibility for this API is guaranteed.

Return type

GraphModule

torch.fx.replace_pattern(gm,pattern,replacement)[source]#

Matches all possible non-overlapping sets of operators and theirdata dependencies (pattern) in the Graph of a GraphModule(gm), then replaces each of these matched subgraphs with anothersubgraph (replacement).

Parameters
Returns

A list ofMatch objects representing the placesin the original graph thatpattern was matched to. The listis empty if there are no matches.Match is defined as:

classMatch(NamedTuple):# Node from which the match was foundanchor:Node# Maps nodes in the pattern subgraph to nodes in the larger graphnodes_map:Dict[Node,Node]

Return type

List[Match]

Examples:

importtorchfromtorch.fximportsymbolic_trace,subgraph_rewriterclassM(torch.nn.Module):def__init__(self)->None:super().__init__()defforward(self,x,w1,w2):m1=torch.cat([w1,w2]).sum()m2=torch.cat([w1,w2]).sum()returnx+torch.max(m1)+torch.max(m2)defpattern(w1,w2):returntorch.cat([w1,w2])defreplacement(w1,w2):returntorch.stack([w1,w2])traced_module=symbolic_trace(M())subgraph_rewriter.replace_pattern(traced_module,pattern,replacement)

The above code will first matchpattern in theforwardmethod oftraced_module. Pattern-matching is done based onuse-def relationships, not node names. For example, if you hadp=torch.cat([a,b]) inpattern, you could matchm=torch.cat([a,b]) in the originalforward function,despite the variable names being different (p vsm).

Thereturn statement inpattern is matched based on itsvalue only; it may or may not match to thereturn statement inthe larger graph. In other words, the pattern doesn’t have to extendto the end of the larger graph.

When the pattern is matched, it will be removed from the largerfunction and replaced byreplacement. If there are multiplematches forpattern in the larger function, each non-overlappingmatch will be replaced. In the case of a match overlap, the firstfound match in the set of overlapping matches will be replaced.(“First” here being defined as the first in a topological orderingof the Nodes’ use-def relationships. In most cases, the first Nodeis the parameter that appears directly afterself, while thelast Node is whatever the function returns.)

One important thing to note is that the parameters of thepattern Callable must be used in the Callable itself,and the parameters of thereplacement Callable must matchthe pattern. The first rule is why, in the above code block, theforward function has parametersx,w1,w2, but thepattern function only has parametersw1,w2.patterndoesn’t usex, so it shouldn’t specifyx as a parameter.As an example of the second rule, consider replacing

defpattern(x,y):returntorch.neg(x)+torch.relu(y)

with

defreplacement(x,y):returntorch.relu(x)

In this case,replacement needs the same number of parametersaspattern (bothx andy), even though the parametery isn’t used inreplacement.

After callingsubgraph_rewriter.replace_pattern, the generatedPython code looks like this:

defforward(self,x,w1,w2):stack_1=torch.stack([w1,w2])sum_1=stack_1.sum()stack_2=torch.stack([w1,w2])sum_2=stack_2.sum()max_1=torch.max(sum_1)add_1=x+max_1max_2=torch.max(sum_2)add_2=add_1+max_2returnadd_2

Note

Backwards-compatibility for this API is guaranteed.

On this page