class torch.fx.Graph [source]
 
Graph is the main data structure used in the FX Intermediate Representation. It consists of a series of Node s, each representing callsites (or other syntactic constructs). The list of Node s, taken together, constitute a valid Python function. For example, the following code import torch
import torch.fx

class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.param = torch.nn.Parameter(torch.rand(3, 4))
        self.linear = torch.nn.Linear(4, 5)

    def forward(self, x):
        return torch.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 : [#users=1] = self.linear.weight
    %add_1 : [#users=1] = call_function[target=operator.add](args = (%x, %linear_weight), kwargs = {})
    %linear_1 : [#users=1] = call_module[target=linear](args = (%add_1,), kwargs = {})
    %relu_1 : [#users=1] = call_method[target=relu](args = (%linear_1,), kwargs = {})
    %sum_1 : [#users=1] = call_function[target=torch.sum](args = (%relu_1,), kwargs = {dim: -1})
    %topk_1 : [#users=1] = call_function[target=torch.topk](args = (%sum_1, 3), kwargs = {})
    return topk_1
 For the semantics of operations represented in the Graph, please see Node.  
__init__() [source]
 
Construct an empty Graph. 
  
call_function(the_function, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_function Node into the Graph. A call_function node represents a call to a Python callable, specified by the_function. the_function can be  Parameters 
 
the_function (Callable[.., Any]) – The function to be called. Can be any PyTorch operator, Python function, or member of the builtins or operator namespaces. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called function. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called function 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.    Returns The newly created and inserted call_function node.  Note The same insertion point and type expression rules apply for this method as Graph.create_node().  
  
call_method(method_name, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_method Node into the Graph. A call_method node represents a call to a given method on the 0th element of args.  Parameters 
 
method_name (str) – The name of the method to apply to the self argument. For example, if args[0] is a Node representing a Tensor, then to call relu() on that Tensor, pass relu to method_name. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called method. Note that this should include a self argument. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called method 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly created and inserted call_method node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node().  
  
call_module(module_name, args=None, kwargs=None, type_expr=None) [source]
 
Insert a call_module Node into the Graph. A call_module node represents a call to the forward() function of a Module in the Module hierarchy.  Parameters 
 
module_name (str) – The qualified name of the Module in the Module hierarchy to be called. For example, if the traced Module has a submodule named foo, which has a submodule named bar, the qualified name foo.bar should be passed as module_name to call that module. 
args (Optional[Tuple[Argument, ..]]) – The positional arguments to be passed to the called method. Note that this should not include a self argument. 
kwargs (Optional[Dict[str, Argument]]) – The keyword arguments to be passed to the called method 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted call_module node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node().  
  
create_node(op, target, args=None, kwargs=None, name=None, type_expr=None) [source]
 
Create a Node and add it to the Graph at the current insert-point. Note that the current insert-point can be set via Graph.inserting_before() and Graph.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 are described in the Graph 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 the Node. This will influence the name of the value assigned to in the Python generated code. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted node.   
  
erase_node(to_erase) [source]
 
Erases a Node from the Graph. Throws an exception if there are still users of that node in the Graph.  Parameters 
to_erase (Node) – The Node to erase from the Graph.   
  
get_attr(qualified_name, type_expr=None) [source]
 
Insert a get_attr node into the Graph. A get_attr Node represents the fetch of an attribute from the Module hierarchy.  Parameters 
 
qualified_name (str) – the fully-qualified name of the attribute to be retrieved. For example, if the traced Module has a submodule named foo, which has a submodule named bar, which has an attribute named baz, the qualified name foo.bar.baz should be passed as qualified_name. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.   Returns 
The newly-created and inserted get_attr node.    Note The same insertion point and type expression rules apply for this method as Graph.create_node.  
  
graph_copy(g, val_map) [source]
 
Copy all nodes from a given graph into self.  Parameters 
 
g (Graph) – The source graph from which to copy Nodes. 
val_map (Dict[Node, Node]) – a dictionary that will be populated with a mapping from nodes in g to nodes in self. Note that val_map can be passed in with values in it already to override copying of certain values.   Returns 
The value in self that is now equivalent to the output value in g, if g had an output node. None otherwise.   
  
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 and then restore it when the with statement exits: with g.inserting_after(n):
    ... # inserting after node n
... # insert point restored to what it was previously
g.inserting_after(n) #  set the insert point permanently
  Parameters 
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__.   
  
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 and then restore it when the with statement exits: with g.inserting_before(n):
    ... # inserting before node n
... # insert point restored to what it was previously
g.inserting_before(n) #  set the insert point permanently
  Parameters 
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__.   
  
lint(root=None) [source]
 
Runs various checks on this Graph to make sure it is well-formed. In particular: - Checks Nodes have correct ownership (owned by this graph) - Checks Nodes appear in topological order - If root is provided, checks that targets exist in root  Parameters 
root (Optional[torch.nn.Module]) – The root module with which to check for targets. This is equivalent to the root argument that is passed when constructing a GraphModule.   
  
node_copy(node, arg_transform=<function Graph.<lambda>>) [source]
 
Copy a node from one graph into another. arg_transform needs to transform arguments from the 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 = {}
for node in g.nodes:
    value_remap[node] = new_graph.node_copy(node, lambda n : value_remap[n])
  Parameters 
 
node (Node) – The node to copy into self. 
arg_transform (Callable[[Node], Argument]) – A function that transforms Node arguments in node’s args and kwargs into the equivalent argument in self. In the simplest case, this should retrieve a value out of a table mapping Nodes in the original graph to self.    
  
property nodes  
Get the list of Nodes that constitute this Graph. Note that this Node list representation is a doubly-linked list. Mutations during iteration (e.g. delete a Node, add a Node) are safe.  Returns 
A doubly-linked list of Nodes. Note that reversed can be called on this list to switch iteration order.   
  
output(result, type_expr=None) [source]
 
Insert an output Node into the Graph. An output node represents a return statement in Python code. result is the value that should be returned.  Parameters 
 
result (Argument) – The value to be returned. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have.     Note The same insertion point and type expression rules apply for this method as Graph.create_node.  
  
placeholder(name, type_expr=None) [source]
 
Insert a placeholder node into the Graph. A placeholder represents a function input.  Parameters 
 
name (str) – A name for the input value. This corresponds to the name of the positional argument to the function this Graph represents. 
type_expr (Optional[Any]) – an optional type annotation representing the Python type the output of this node will have. This is needed in some cases for proper code generation (e.g. when the function is used subsequently in TorchScript compilation).     Note The same insertion point and type expression rules apply for this method as Graph.create_node.  
  
print_tabular() [source]
 
Prints the intermediate representation of the graph in tabular format. 
  
python_code(root_module) [source]
 
Turn this Graph into valid Python code.  Parameters 
root_module (str) – The name of the root module on which to look-up qualified name targets. This is usually ‘self’.  Returns 
The string source code generated from this Graph.