Casts this storage to bfloat16 type
Casts this storage to bool type
Casts this storage to byte type
Casts this storage to char type
Returns a copy of this storage
Casts this storage to complex double type
Casts this storage to complex float type
Returns a CPU copy of this storage if it’s not already on the CPU
Returns a copy of this object in CUDA memory.
If this object is already in CUDA memory and on the correct device, then no copy is performed and the original object is returned.
Casts this storage to double type
Casts this storage to float type
If shared is True, then memory is shared between all processes. All changes are written to the file. If shared is False, then the changes on the storage do not affect the file.
size is the number of elements in the storage. If shared is False, then the file must contain at least size * sizeof(Type) bytes (Type is the type of storage). If shared is True the file will be created if needed.
Casts this storage to half type
Casts this storage to int type
Casts this storage to long type
Copies the storage to pinned memory, if it’s not already pinned.
Moves the storage to shared memory.
This is a no-op for storages already in shared memory and for CUDA storages, which do not need to be moved for sharing across processes. Storages in shared memory cannot be resized.
Returns: self
Casts this storage to short type
Returns a list containing the elements of this storage
Returns the type if dtype is not provided, else casts this object to the specified type.
If this is already of the correct type, no copy is performed and the original object is returned.


Casts this storage to complex double type


Casts this storage to complex float type


Returns a CPU copy of this storage if it’s not already on the CPU


Returns a copy of this object in CUDA memory.
If this object is already in CUDA memory and on the correct device, then no copy is performed and the original object is returned.


If shared is True, then memory is shared between all processes. All changes are written to the file. If shared is False, then the changes on the storage do not affect the file.
size is the number of elements in the storage. If shared is False, then the file must contain at least size * sizeof(Type) bytes (Type is the type of storage). If shared is True the file will be created if needed.


Copies the storage to pinned memory, if it’s not already pinned.


Moves the storage to shared memory.
This is a no-op for storages already in shared memory and for CUDA storages, which do not need to be moved for sharing across processes. Storages in shared memory cannot be resized.
Returns: self


Returns a list containing the elements of this storage


Returns the type if dtype is not provided, else casts this object to the specified type.
If this is already of the correct type, no copy is performed and the original object is returned.


Raises input to the power of exponent, elementwise, in double precision. If neither input is complex returns a torch.float64 tensor, and if one or more inputs is complex returns a torch.complex128 tensor.
Note
This function always computes in double precision, unlike torch.pow(), which implements more typical type promotion. This is useful when the computation needs to be performed in a wider or more precise dtype, or the results of the computation may contain fractional values not representable in the input dtypes, like when an integer base is raised to a negative integer exponent.
out (Tensor, optional) – the output tensor.
Example:


Returns a new tensor with the floor of the elements of input, the largest integer less than or equal to each element.
input (Tensor) – the input tensor.
out (Tensor, optional) – the output tensor.
Example:


Warning
This function’s name is a misnomer. It actually rounds the quotient towards zero instead of taking its floor. This behavior will be deprecated in a future PyTorch release.
Computes input divided by other, elementwise, and rounds each quotient towards zero. Equivalently, it truncates the quotient(s):
Supports broadcasting to a common shape, type promotion, and integer and float inputs.
out (Tensor, optional) – the output tensor.
Example:


Computes the element-wise maximum of input and other.
This is like torch.maximum() except it handles NaNs differently: if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the maximum. Only if both elements are NaN is NaN propagated.
This function is a wrapper around C++’s std::fmax and is similar to NumPy’s fmax function.
Supports broadcasting to a common shape, type promotion, and integer and floating-point inputs.
out (Tensor, optional) – the output tensor.
Example:


Computes the element-wise minimum of input and other.
This is like torch.minimum() except it handles NaNs differently: if exactly one of the two elements being compared is a NaN then the non-NaN element is taken as the minimum. Only if both elements are NaN is NaN propagated.
This function is a wrapper around C++’s std::fmin and is similar to NumPy’s fmin function.
Supports broadcasting to a common shape, type promotion, and integer and floating-point inputs.
out (Tensor, optional) – the output tensor.
Example:


Computes the element-wise remainder of division.
The dividend and divisor may contain both for integer and floating point numbers. The remainder has the same sign as the dividend input.
Supports broadcasting to a common shape, type promotion, and integer and float inputs.
Note
When the divisor is zero, returns NaN for floating point dtypes on both CPU and GPU; raises RuntimeError for integer division by zero on CPU; Integer division by zero on GPU may return any value.
out (Tensor, optional) – the output tensor.
Example:


Computes the fractional portion of each element in input.
Example:


Creates a Tensor from a numpy.ndarray.
The returned tensor and ndarray share the same memory. Modifications to the tensor will be reflected in the ndarray and vice versa. The returned tensor is not resizable.
It currently accepts ndarray with dtypes of numpy.float64, numpy.float32, numpy.float16, numpy.complex64, numpy.complex128, numpy.int64, numpy.int32, numpy.int16, numpy.int8, numpy.uint8, and numpy.bool.
Example:


Creates a tensor of size size filled with fill_value. The tensor’s dtype is inferred from fill_value.
Example:


Returns a tensor with the same size as input filled with fill_value. torch.full_like(input, fill_value) is equivalent to torch.full(input.size(), fill_value, dtype=input.dtype, layout=input.layout, device=input.device).


Warning
The torch.futures package is experimental and subject to change.
This package provides a Future type that encapsulates an asynchronous execution and a set of utility functions to simplify operations on Future objects. Currently, the Future type is primarily used by the Distributed RPC Framework.
Wrapper around a torch._C.Future which encapsulates an asynchronous execution of a callable, e.g. rpc_async(). It also exposes a set of APIs to add callback functions and set results.
Return True if this Future is done. A Future is done if it has a result or an exception.
Set an exception for this Future, which will mark this Future as completed with an error and trigger all attached callbacks. Note that when calling wait()/value() on this Future, the exception set here will be raised inline.
result (BaseException) – the exception for this Future.
Set the result for this Future, which will mark this Future as completed and trigger all attached callbacks. Note that a Future cannot be marked completed twice.
result (object) – the result object of this Future.
Append the given callback function to this Future, which will be run when the Future is completed. Multiple callbacks can be added to the same Future, and will be invoked in the same order as they were added. The callback must take one argument, which is the reference to this Future. The callback function can use the Future.wait() API to get the value. Note that if this Future is already completed, the given callback will be run immediately inline.
callback (Callable) – a Callable that takes this Future as the only argument.
A new Future object that holds the return value of the callback and will be marked as completed when the given callback finishes.
Block until the value of this Future is ready.
The value held by this Future. If the function (callback or RPC) creating the value has thrown an error, this wait method will also throw an error.
Collects the provided Future objects into a single combined Future that is completed when all of the sub-futures are completed.
futures (list) – a list of Future objects.
Returns a Future object to a list of the passed in Futures.
Waits for all provided futures to be complete, and returns the list of completed values.
futures (list) – a list of Future object.
A list of the completed Future results. This method will throw an error if wait on any Future throws.


Collects the provided Future objects into a single combined Future that is completed when all of the sub-futures are completed.
futures (list) – a list of Future objects.
Returns a Future object to a list of the passed in Futures.


Wrapper around a torch._C.Future which encapsulates an asynchronous execution of a callable, e.g. rpc_async(). It also exposes a set of APIs to add callback functions and set results.
Return True if this Future is done. A Future is done if it has a result or an exception.
Set an exception for this Future, which will mark this Future as completed with an error and trigger all attached callbacks. Note that when calling wait()/value() on this Future, the exception set here will be raised inline.
result (BaseException) – the exception for this Future.
Set the result for this Future, which will mark this Future as completed and trigger all attached callbacks. Note that a Future cannot be marked completed twice.
result (object) – the result object of this Future.
Append the given callback function to this Future, which will be run when the Future is completed. Multiple callbacks can be added to the same Future, and will be invoked in the same order as they were added. The callback must take one argument, which is the reference to this Future. The callback function can use the Future.wait() API to get the value. Note that if this Future is already completed, the given callback will be run immediately inline.
callback (Callable) – a Callable that takes this Future as the only argument.
A new Future object that holds the return value of the callback and will be marked as completed when the given callback finishes.
Block until the value of this Future is ready.
The value held by this Future. If the function (callback or RPC) creating the value has thrown an error, this wait method will also throw an error.


Return True if this Future is done. A Future is done if it has a result or an exception.


Set an exception for this Future, which will mark this Future as completed with an error and trigger all attached callbacks. Note that when calling wait()/value() on this Future, the exception set here will be raised inline.
result (BaseException) – the exception for this Future.


Set the result for this Future, which will mark this Future as completed and trigger all attached callbacks. Note that a Future cannot be marked completed twice.
result (object) – the result object of this Future.


Append the given callback function to this Future, which will be run when the Future is completed. Multiple callbacks can be added to the same Future, and will be invoked in the same order as they were added. The callback must take one argument, which is the reference to this Future. The callback function can use the Future.wait() API to get the value. Note that if this Future is already completed, the given callback will be run immediately inline.
callback (Callable) – a Callable that takes this Future as the only argument.
A new Future object that holds the return value of the callback and will be marked as completed when the given callback finishes.


Block until the value of this Future is ready.
The value held by this Future. If the function (callback or RPC) creating the value has thrown an error, this wait method will also throw an error.


Waits for all provided futures to be complete, and returns the list of completed values.
futures (list) – a list of Future object.
A list of the completed Future results. This method will throw an error if wait on any Future throws.


This feature is under a Beta release and its API may change.
FX is a toolkit for developers to use to transform nn.Module instances. FX consists of three main components: a symbolic tracer, an intermediate representation, and Python code generation. A demonstration of these components in action:
The symbolic tracer performs “symbolic execution” of the Python code. It feeds fake values, called Proxies, through the code. Operations on theses Proxies are recorded. More information about symbolic tracing can be found in the symbolic_trace() and Tracer documentation.
The intermediate representation is the container for the operations that were recorded during symbolic tracing. It consists of a list of Nodes that represent function inputs, callsites (to functions, methods, or torch.nn.Module instances), and return values. More information about the IR can be found in the documentation for Graph. The IR is the format on which transformations are applied.
Python code generation is what makes FX a Python-to-Python (or Module-to-Module) transformation toolkit. For each Graph IR, we can create valid Python code matching the Graph’s semantics. This functionality is wrapped up in GraphModule, which is a torch.nn.Module instance that holds a Graph as well as a forward 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. In addition, these components can be used separately. For example, symbolic tracing can be used in isolation to capture a form of the code for analysis (and not transformation) purposes. Code generation can be used for programmatically generating models, for example from a config file. There are many uses for FX!
Several example transformations can be found at the examples repository.
What is an FX transform? Essentially, it’s a function that looks like this.
Your transform will take in an torch.nn.Module, acquire a Graph from it, do some modifications, and return a new torch.nn.Module. You should think of the torch.nn.Module that your FX transform returns as identical to a regular torch.nn.Module – you can pass it to another FX transform, you can pass it to TorchScript, or you can run it. Ensuring that the inputs and outputs of your FX transform are a torch.nn.Module will allow for composability.
Note
It is also possible to modify an existing GraphModule instead of creating a new one, like so:
Note that you MUST call GraphModule.recompile() to bring the generated forward() method on the GraphModule in sync with the modified Graph.
Given that you’ve passed in a torch.nn.Module that has been traced into a Graph, there are now two primary approaches you can take to building a new Graph.
Full treatment of the semantics of graphs can be found in the Graph documentation, but we are going to cover the basics here. A Graph is a data structure that represents a method on a GraphModule. The information that this requires is:
All three of these concepts are represented with Node instances. Let’s see what we mean by that with a short example:
Here we define a module MyModule for demonstration purposes, instantiate it, symbolically trace it, then call the Graph.print_tabular() method to print out a table showing the nodes of this Graph:
opcode
name
target
args
kwargs
placeholder
x
x
()
{}
get_attr
linear_weight
linear.weight
()
{}
call_function
add_1
<built-in function add>
(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.
Given that we now know the basics of how code is represented in FX, we can now explore how we would edit a Graph.
One approach to building this new Graph is to directly manipulate your old one. To aid in this, we can simply take the Graph we obtain from symbolic tracing and modify it. For example, let’s say we desire to replace torch.add() calls with torch.mul() calls.
We can also do more involved Graph rewrites, such as deleting or appending nodes. To aid in these transformations, FX has utility functions for transforming the graph that can be found in the Graph documentation. An example of using these APIs to append a torch.relu() call can be found below.
For simple transformations that only consist of substitutions, you can also make use of the subgraph rewriter.
FX also provides another level of automation on top of direct graph manipulation. The replace_pattern() API is essentially a “find/replace” tool for editing Graphs. It allows you to specify a pattern and replacement function and it will trace through those functions, find instances of the group of operations in the pattern graph, and replace those instances with copies of the replacement graph. This can help to greatly automate tedious graph manipulation code, which can get unwieldy as the transformations get more complex.
Another way of manipulating Graphs is by reusing the Proxy machinery used in symbolic tracing. For example, let’s imagine that we wanted to write a transformation that decomposed PyTorch functions into smaller operations. It would transform every F.relu(x) call into (x > 0) * x. One possibility would be to perform the requisite graph rewriting to insert the comparison and multiplication after the F.relu, and then clean up the original F.relu. However, we can automate this process by using Proxy objects to automatically record operations into the Graph.
To use this method, we write the operations that we want inserted as regular PyTorch code and invoke that code with Proxy objects as arugments. These Proxy objects will capture the operations that are performed on them and append them to the Graph.
In addition to avoiding explicit graph manipulation, using Proxys also 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 and maintainability of the rules.
A worked example of using Proxys for Graph manipulation can be found here.
A useful code organizational pattern in FX is to loop over all the Nodes in a Graph and execute them. This can be used for several things including runtime analysis of values flowing through the graph or transformation of the code via retracing with Proxys. For example, suppose we want to run a GraphModule and record the torch.Tensor shape and dtype properties on the nodes as we see them at runtime. That might look like:
As you can see, a full interpreter for FX is not that complicated but it can be very useful. To ease using this pattern, we provide the Interpreter class, which encompasses the above logic in a way that certain aspects of the interpreter’s execution can be overridden via method overrides.
In addition to executing operations, we can also generate a new Graph by feeding Proxy values through an interpreter. Similarly, we provide the Transformer class to encompass this pattern. Transformer behaves similarly to Interpreter, but instead of calling the run method to get a concrete output value from the Module, you would call the Transformer.transform() method to return a new GraphModule which was subject to any transformation rules you installed as overridden methods.
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 work backwards: first, check the results of invoking the generated module to prove or disprove correctness. Then, inspect and debug the generated code. Then, debug the process of transformations that led to the generated code.
If you’re not familiar with debuggers, please see the auxiliary section Available Debuggers.
Because the output of most deep learning modules consists of floating point torch.Tensor instances, checking for equivalence between the results of two torch.nn.Module is not as straightforward as doing a simple equality check. To motivate this, let’s use an example:
Here, we’ve tried to check equality of the values of two deep learning models with the == equality operator. However, this is not well- defined both due to the issue of that operator returning a tensor and not a bool, but also because comparison of floating point values should use a margin of error (or epsilon) to account for the non-commutativity of floating point operations (see here for more details). We can use torch.allclose() instead, which will give us an approximate comparison taking into account a relative and absolute tolerance threshold:
This is the first tool in our toolbox to check if transformed modules are behaving as we expect compared to a reference implementation.
Because FX generates the forward() function on GraphModules, using traditional debugging techniques like print statements or pdb is not as straightfoward. Luckily, we have several techniques we can use for debugging the generated code.
Invoke pdb to step into the running program. Although the code that represents the Graph is not in any source file, we can still step into it manually using pdb when the forward pass is invoked.
If you’d like to run the same code multiple times, then it can be a bit tedious to step to the right code with pdb. In that case, one approach is to simply copy-paste the generated forward pass into your code and examine it from there.
GraphModule.to_folder() is a method in GraphModule that allows you to dump out the generated FX code to a folder. Although copying the forward pass into the code often suffices as in Print the Generated Code, it may be easier to examine modules and parameters using to_folder.
After running the above example, we can then look at the code within foo/module.py and modify it as desired (e.g. adding print statements or using pdb) to debug the generated code.
Now that we’ve identified that a transformation is creating incorrect code, it’s time to debug the transformation itself. First, we’ll check the Limitations of Symbolic Tracing section in the documentation. Once we verify that tracing is working as expected, the goal becomes figuring out what went wrong during our GraphModule transformation. There may be a quick answer in Writing Transformations, but, if not, there are several ways to examine our traced module:
Using the utility functions above, we can compare our traced Module before and after we’ve applied our transformations. Sometimes, a simple visual comparison is enough to trace down a bug. If it’s still not clear what’s going wrong, a debugger like pdb can be a good next step.
Going off of the example above, consider the following code:
Using the above example, let’s say that the call to print(traced) showed us that there was an error in our transforms. We want to find what goes wrong using a debugger. We start a pdb session. We can see what’s happening during the transform by breaking on transform_graph(traced), then pressing s to “step into” the call to transform_graph(traced).
We may also have good luck by editing the print_tabular method to print different attributes of the Nodes in the Graph. (For example, we might want to see the Node’s input_nodes and users.)
The most common Python debugger is pdb. You can start your program in “debug mode” with pdb by typing python -m pdb FILENAME.py into the command line, where FILENAME is the name of the file you want to debug. After that, you can use the pdb debugger commands to move through your running program stepwise. It’s common to set a breakpoint (b LINE-NUMBER) when you start pdb, then call c to run the program until that point. This prevents you from having to step through each line of execution (using s or n) to get to the part of the code you want to examine. Alternatively, you can write import pdb; pdb.set_trace() before the line you want to break at. If you add pdb.set_trace(), your program will automatically start in debug mode when you run it. (In other words, you can just type python FILENAME.py into the command line instead of python -m pdb FILENAME.py.) Once you’re running your file in debug mode, you can step through the code and examine your program’s internal state using certain commands. There are many excellent tutorials on pdb online, including RealPython’s “Python Debugging With Pdb”.
IDEs like PyCharm or VSCode usually have a debugger built in. In your IDE, you can choose to either a) use pdb by pulling up a terminal window in your IDE (e.g. View → Terminal in VSCode), or b) use the built-in debugger (usually a graphical wrapper around pdb).
FX uses a system of symbolic tracing (a.k.a symbolic execution) to capture the semantics of programs in a transformable/analyzable form. The system is tracing in that it executes the program (really a torch.nn.Module or function) to record operations. It is symbolic in that the data flowing through the program during this execution is not real data, but rather symbols (Proxy in FX parlance).
Although symbolic tracing works for most neural net code, it has some limitations.
The main limitation of symbolic tracing is it does not currently support dynamic control flow. That is, loops or if statements where the condition may depend on the input values of the program.
For example, let’s examine the following program:
The condition to the if statement relies on the value of dim0, which eventually relies on the value of x, a function input. Since x can change (i.e. if you pass a new input tensor to the traced function), this is dynamic control flow. The traceback walks back up through your code to show you where this situation happens.
On the other hand, so-called static control flow is supported. Static control flow is loops or if statements whose value cannot change across invocations. Typically, in PyTorch programs, this control flow arises for code making decisions about a model’s architecture based on hyper-parameters. As a concrete example:
The if-statement if self.do_activation does not depend on any function inputs, thus it is static. do_activation can be considered to be a hyper-parameter, and the traces of different instances of MyModule with different values for that parameter have different code. This is a valid pattern that is supported by symbolic tracing.
Many instances of dynamic control flow are semantically static control flow. These instances can be made to support symbolic tracing by removing the data dependencies on input values, for example by moving values to Module attributes or by passing constant values during symbolic tracing:
In the case of truly dynamic control flow, the sections of the program that contain this code can be traced as calls to the Method (see Customizing Tracing with the Tracer class) or function (see wrap()) rather than tracing through them.
FX uses __torch_function__ as the mechanism by which it intercepts calls (see the technical overview for more information about this). Some functions, such as builtin Python functions or those in the math module, are things that are not covered by __torch_function__, but we would still like to capture them in symbolic tracing. For example:
The error tells us that the built-in function len is not supported. We can make it so that functions like this are recorded in the trace as direct calls using the wrap() API:
The Tracer class is the class that underlies the implementation of symbolic_trace. The behavior of tracing can be customized by subclassing Tracer, like so:
Leaf Modules are the modules that appear as calls in the symbolic trace rather than being traced through. The default set of leaf modules is the set of standard torch.nn module instances. For example:
The set of leaf modules can be customized by overriding Tracer.is_leaf_module().
Tensor constructors (e.g. torch.zeros, torch.ones, torch.rand, torch.randn, torch.sparse_coo_tensor) are currently not traceable.
Type annotations
Symbolic tracing API
Given an nn.Module or function instance root, this function will return a GraphModule constructed by recording operations seen while tracing through root.
a Module created from the recorded operations from root.
GraphModule
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 being traced through:
This function can also equivalently be used as a decorator:
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 trace rather than traced through.
fn_or_name (Union[str, Callable]) – The function or name of the global function to insert into the graph when it’s called
GraphModule is an nn.Module generated from an fx.Graph. Graphmodule has a graph attribute, as well as code and forward attributes generated from that graph.
Warning
When graph is reassigned, code and forward will be automatically regenerated. However, if you edit the contents of the graph without reassigning the graph attribute itself, you must call recompile() to update the generated code.
Construct a GraphModule.
Return the Python code generated from the Graph underlying this GraphModule.
Return the Graph underlying this GraphModule
Recompile this GraphModule from its graph attribute. This should be called after editing the contained graph, otherwise the generated code of this GraphModule will be out of date.
Dumps out module to folder with module_name so that it can be imported with from <folder> import <module_name>
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
Will produce the following Graph:
For the semantics of operations represented in the Graph, please see Node.
Construct an empty Graph.
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
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().
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.
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().
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.
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 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().
The newly-created and inserted node.
Erases a Node from the Graph. Throws an exception if there are still users of that node in the Graph.
to_erase (Node) – The Node to erase from the Graph.
Insert a get_attr node into the Graph. A get_attr Node represents the fetch of an attribute from the Module hierarchy.
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.
Copy all nodes from a given graph into self.
The value in self that is now equivalent to the output value in g, if g had an output node. None otherwise.
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:
n (Optional[Node]) – The node before which to insert. If None this will insert after the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.
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:
n (Optional[Node]) – The node before which to insert. If None this will insert before the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.
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
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.
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:
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.
A doubly-linked list of Nodes. Note that reversed can be called on this list to switch iteration order.
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.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.
Insert a placeholder node into the Graph. A placeholder represents a function input.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.
Prints the intermediate representation of the graph in tabular format.
Turn this Graph into valid Python code.
root_module (str) – The name of the root module on which to look-up qualified name targets. This is usually ‘self’.
The string source code generated from this Graph.
Node is the data structure that represents individual operations within a Graph. For the most part, Nodes represent callsites to various entities, such as operators, methods, and Modules (some exceptions include nodes that specify function inputs and outputs). Each Node has a function specified by its op property. The Node semantics for each value of op are as follows:
Return all Nodes that are inputs to this Node. This is equivalent to iterating over args and kwargs and only collecting the values that are Nodes.
List of Nodes that appear in the args and kwargs of this Node, in that order.
Insert x after this node in the list of nodes in the graph. Equvalent to self.next.prepend(x)
x (Node) – The node to put after this node. Must be a member of the same graph.
The tuple of arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
The dict of keyword arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
Returns the next Node in the linked list of Nodes.
The next Node in the linked list of Nodes.
Insert x before this node in the list of nodes in the graph. Example:
x (Node) – The node to put before this node. Must be a member of the same graph.
Returns the previous Node in the linked list of Nodes.
The previous Node in the linked list of Nodes.
Replace all uses of self in the Graph with the Node replace_with.
replace_with (Node) – The node to replace all uses of self with.
The list of Nodes on which this change was made.
Tracer is the class that implements the symbolic tracing functionality of torch.fx.symbolic_trace. A call to symbolic_trace(m) is equivalent to Tracer().trace(m).
Tracer can be subclassed to override various behaviors of the tracing process. The different behaviors that can be overridden are described in the docstrings of the methods on this class.
Method that specifies the behavior of this Tracer when it encounters a call to an nn.Module instance.
By default, the behavior is to check if the called module is a leaf module via is_leaf_module. If it is, emit a call_module node referring to m in the Graph. Otherwise, call the Module normally, tracing through the operations in its forward function.
This method can be overridden to–for example–create nested traced GraphModules, or any other behavior you would want while tracing across Module boundaries. Module boundaries.
The return value from the Module call. In the case that a call_module node was emitted, this is a Proxy value. Otherwise, it is whatever value was returned from the Module invocation.
A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the Graph.
By default, the behavior includes:
Given a non-Proxy Tensor object, emit IR for various cases:
This method can be overridden to support more types.
a (Any) – The value to be emitted as an Argument in the Graph.
The value a converted into the appropriate Argument
Create placeholder nodes corresponding to the signature of the root Module. This method introspects root’s signature and emits those nodes accordingly, also supporting *args and **kwargs.
A method to specify whether a given nn.Module is a “leaf” module.
Leaf modules are the atomic units that appear in the IR, referenced by call_module calls. By default, Modules in the PyTorch standard library namespace (torch.nn) are leaf modules. All other modules are traced through and their constituent ops are recorded, unless specified otherwise via this parameter.
Helper method to find the qualified name of mod in the Module hierarchy of root. For example, if root has a submodule named foo, which has a submodule named bar, passing bar into this function will return the string “foo.bar”.
mod (str) – The Module to retrieve the qualified name for.
Trace root and return the corresponding FX Graph representation. root can either be an nn.Module instance or a Python callable.
Note that after this call, self.root may be different from the root passed in here. For example, when a free function is passed to trace(), we will create an nn.Module instance to use as the root and add embedded constants to.
root (Union[Module, Callable]) – Either a Module or a function to be traced through.
A Graph representing the semantics of the passed-in root.
Proxy objects are Node wrappers that flow through the program during symbolic tracing and record all the operations (torch function calls, method calls, operators) that they touch into the growing FX Graph.
If you’re doing graph transforms, you can wrap your own Proxy method around a raw Node so that you can use the overloaded operators to add additional things to a Graph.
An Interpreter executes an FX graph Node-by-Node. This pattern can be useful for many things, including writing code transformations as well as analysis passes.
Methods in the Interpreter class can be overridden to customize the behavior of execution. The map of overrideable methods in terms of call hierarchy:
Suppose we want to swap all instances of torch.neg with torch.sigmoid and vice versa (including their Tensor method equivalents). We could subclass Interpreter like so:
module (GraphModule) – The module to be executed
Execute a call_function node and return the result.
Any: The value returned by the function invocation
Execute a call_method node and return the result.
Any: The value returned by the method invocation
Execute a call_module node and return the result.
Any: The value returned by the module invocation
Fetch the concrete values of args and kwargs of node n from the current execution environment.
n (Node) – The node for which args and kwargs should be fetched.
args and kwargs with concrete values for n.
Tuple[Tuple, Dict]
Fetch an attribute from the Module hierarchy of self.module.
target (str) – The fully-qualfiied name of the attribute to fetch
The value of the attribute.
Any
Execute a get_attr node. Will retrieve an attribute value from the Module hierarchy of self.module.
The value of the attribute that was retrieved
Any
Recursively descend through args and look up the concrete value for each Node in the current execution environment.
Execute an output node. This really just retrieves the value referenced by the output node and returns it.
The return value referenced by the output node
Any
Execute a placeholder node. Note that this is stateful: Interpreter maintains an internal iterator over arguments passed to run and this method returns next() on that iterator.
The argument value that was retrieved.
Any
Run module via interpretation and return the result.
The value returned from executing the Module
Any
Run a specific node n and return the result. Calls into placeholder, get_attr, call_function, call_method, call_module, or output depending on node.op
n (Node) – The Node to execute
The result of executing n
Any
Transformer is a special type of interpreter that produces a new Module. It exposes a transform() method that returns the transformed Module. Transformer does not require arguments to run, as Interpreter does. Transformer works entirely symbolically.
Suppose we want to swap all instances of torch.neg with torch.sigmoid and vice versa (including their Tensor method equivalents). We could subclass Transformer like so:
module (GraphModule) – The Module to be transformed.
Execute a get_attr node. In Transformer, this is overridden to insert a new get_attr node into the output graph.
Execute a placeholder node. In Transformer, this is overridden to insert a new placeholder into the output graph.
Transform self.module and return the transformed GraphModule.
Matches all possible non-overlapping sets of operators and their data dependencies (pattern) in the Graph of a GraphModule (gm), then replaces each of these matched subgraphs with another subgraph (replacement).
A list of Match objects representing the places in the original graph that pattern was matched to. The list is empty if there are no matches. Match is defined as:
List[Match]
Examples:
The above code will first match pattern in the forward method of traced_module. Pattern-matching is done based on use-def relationships, not node names. For example, if you had p = torch.cat([a, b]) in pattern, you could match m = torch.cat([a, b]) in the original forward function, despite the variable names being different (p vs m).
The return statement in pattern is matched based on its value only; it may or may not match to the return statement in the larger graph. In other words, the pattern doesn’t have to extend to the end of the larger graph.
When the pattern is matched, it will be removed from the larger function and replaced by replacement. If there are multiple matches for pattern in the larger function, each non-overlapping match will be replaced. In the case of a match overlap, the first found match in the set of overlapping matches will be replaced. (“First” here being defined as the first in a topological ordering of the Nodes’ use-def relationships. In most cases, the first Node is the parameter that appears directly after self, while the last Node is whatever the function returns.)
One important thing to note is that the parameters of the pattern Callable must be used in the Callable itself, and the parameters of the replacement Callable must match the pattern. The first rule is why, in the above code block, the forward function has parameters x, w1, w2, but the pattern function only has parameters w1, w2. pattern doesn’t use x, so it shouldn’t specify x as a parameter. As an example of the second rule, consider replacing
with
In this case, replacement needs the same number of parameters as pattern (both x and y), even though the parameter y isn’t used in replacement.
After calling subgraph_rewriter.replace_pattern, the generated Python code looks like this:


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
Will produce the following Graph:
For the semantics of operations represented in the Graph, please see Node.
Construct an empty Graph.
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
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().
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.
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().
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.
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 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().
The newly-created and inserted node.
Erases a Node from the Graph. Throws an exception if there are still users of that node in the Graph.
to_erase (Node) – The Node to erase from the Graph.
Insert a get_attr node into the Graph. A get_attr Node represents the fetch of an attribute from the Module hierarchy.
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.
Copy all nodes from a given graph into self.
The value in self that is now equivalent to the output value in g, if g had an output node. None otherwise.
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:
n (Optional[Node]) – The node before which to insert. If None this will insert after the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.
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:
n (Optional[Node]) – The node before which to insert. If None this will insert before the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.
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
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.
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:
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.
A doubly-linked list of Nodes. Note that reversed can be called on this list to switch iteration order.
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.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.
Insert a placeholder node into the Graph. A placeholder represents a function input.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.
Prints the intermediate representation of the graph in tabular format.
Turn this Graph into valid Python code.
root_module (str) – The name of the root module on which to look-up qualified name targets. This is usually ‘self’.
The string source code generated from this Graph.


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
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().


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.
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().


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.
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 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().
The newly-created and inserted node.


Erases a Node from the Graph. Throws an exception if there are still users of that node in the Graph.
to_erase (Node) – The Node to erase from the Graph.


Insert a get_attr node into the Graph. A get_attr Node represents the fetch of an attribute from the Module hierarchy.
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.


Copy all nodes from a given graph into self.
The value in self that is now equivalent to the output value in g, if g had an output node. None otherwise.


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:
n (Optional[Node]) – The node before which to insert. If None this will insert after the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.


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:
n (Optional[Node]) – The node before which to insert. If None this will insert before the beginning of the entire graph.
A resource manager that will restore the insert point on __exit__.


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
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.


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.
A doubly-linked list of Nodes. Note that reversed can be called on this list to switch iteration order.


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:


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.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.


Insert a placeholder node into the Graph. A placeholder represents a function input.
Note
The same insertion point and type expression rules apply for this method as Graph.create_node.


Prints the intermediate representation of the graph in tabular format.


Turn this Graph into valid Python code.
root_module (str) – The name of the root module on which to look-up qualified name targets. This is usually ‘self’.
The string source code generated from this Graph.


GraphModule is an nn.Module generated from an fx.Graph. Graphmodule has a graph attribute, as well as code and forward attributes generated from that graph.
Warning
When graph is reassigned, code and forward will be automatically regenerated. However, if you edit the contents of the graph without reassigning the graph attribute itself, you must call recompile() to update the generated code.
Construct a GraphModule.
Return the Python code generated from the Graph underlying this GraphModule.
Return the Graph underlying this GraphModule
Recompile this GraphModule from its graph attribute. This should be called after editing the contained graph, otherwise the generated code of this GraphModule will be out of date.
Dumps out module to folder with module_name so that it can be imported with from <folder> import <module_name>


Return the Python code generated from the Graph underlying this GraphModule.


Return the Graph underlying this GraphModule


Recompile this GraphModule from its graph attribute. This should be called after editing the contained graph, otherwise the generated code of this GraphModule will be out of date.


Dumps out module to folder with module_name so that it can be imported with from <folder> import <module_name>


An Interpreter executes an FX graph Node-by-Node. This pattern can be useful for many things, including writing code transformations as well as analysis passes.
Methods in the Interpreter class can be overridden to customize the behavior of execution. The map of overrideable methods in terms of call hierarchy:
Suppose we want to swap all instances of torch.neg with torch.sigmoid and vice versa (including their Tensor method equivalents). We could subclass Interpreter like so:
module (GraphModule) – The module to be executed
Execute a call_function node and return the result.
Any: The value returned by the function invocation
Execute a call_method node and return the result.
Any: The value returned by the method invocation
Execute a call_module node and return the result.
Any: The value returned by the module invocation
Fetch the concrete values of args and kwargs of node n from the current execution environment.
n (Node) – The node for which args and kwargs should be fetched.
args and kwargs with concrete values for n.
Tuple[Tuple, Dict]
Fetch an attribute from the Module hierarchy of self.module.
target (str) – The fully-qualfiied name of the attribute to fetch
The value of the attribute.
Any
Execute a get_attr node. Will retrieve an attribute value from the Module hierarchy of self.module.
The value of the attribute that was retrieved
Any
Recursively descend through args and look up the concrete value for each Node in the current execution environment.
Execute an output node. This really just retrieves the value referenced by the output node and returns it.
The return value referenced by the output node
Any
Execute a placeholder node. Note that this is stateful: Interpreter maintains an internal iterator over arguments passed to run and this method returns next() on that iterator.
The argument value that was retrieved.
Any
Run module via interpretation and return the result.
The value returned from executing the Module
Any
Run a specific node n and return the result. Calls into placeholder, get_attr, call_function, call_method, call_module, or output depending on node.op
n (Node) – The Node to execute
The result of executing n
Any


Execute a call_function node and return the result.
Any: The value returned by the function invocation


Execute a call_method node and return the result.
Any: The value returned by the method invocation


Execute a call_module node and return the result.
Any: The value returned by the module invocation


Fetch the concrete values of args and kwargs of node n from the current execution environment.
n (Node) – The node for which args and kwargs should be fetched.
args and kwargs with concrete values for n.
Tuple[Tuple, Dict]


Fetch an attribute from the Module hierarchy of self.module.
target (str) – The fully-qualfiied name of the attribute to fetch
The value of the attribute.
Any


Execute a get_attr node. Will retrieve an attribute value from the Module hierarchy of self.module.
The value of the attribute that was retrieved
Any


Recursively descend through args and look up the concrete value for each Node in the current execution environment.


Execute an output node. This really just retrieves the value referenced by the output node and returns it.
The return value referenced by the output node
Any


Execute a placeholder node. Note that this is stateful: Interpreter maintains an internal iterator over arguments passed to run and this method returns next() on that iterator.
The argument value that was retrieved.
Any


Run module via interpretation and return the result.
The value returned from executing the Module
Any


Run a specific node n and return the result. Calls into placeholder, get_attr, call_function, call_method, call_module, or output depending on node.op
n (Node) – The Node to execute
The result of executing n
Any


Node is the data structure that represents individual operations within a Graph. For the most part, Nodes represent callsites to various entities, such as operators, methods, and Modules (some exceptions include nodes that specify function inputs and outputs). Each Node has a function specified by its op property. The Node semantics for each value of op are as follows:
Return all Nodes that are inputs to this Node. This is equivalent to iterating over args and kwargs and only collecting the values that are Nodes.
List of Nodes that appear in the args and kwargs of this Node, in that order.
Insert x after this node in the list of nodes in the graph. Equvalent to self.next.prepend(x)
x (Node) – The node to put after this node. Must be a member of the same graph.
The tuple of arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
The dict of keyword arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.
Returns the next Node in the linked list of Nodes.
The next Node in the linked list of Nodes.
Insert x before this node in the list of nodes in the graph. Example:
x (Node) – The node to put before this node. Must be a member of the same graph.
Returns the previous Node in the linked list of Nodes.
The previous Node in the linked list of Nodes.
Replace all uses of self in the Graph with the Node replace_with.
replace_with (Node) – The node to replace all uses of self with.
The list of Nodes on which this change was made.


Return all Nodes that are inputs to this Node. This is equivalent to iterating over args and kwargs and only collecting the values that are Nodes.
List of Nodes that appear in the args and kwargs of this Node, in that order.


Insert x after this node in the list of nodes in the graph. Equvalent to self.next.prepend(x)
x (Node) – The node to put after this node. Must be a member of the same graph.


The tuple of arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.


The dict of keyword arguments to this Node. The interpretation of arguments depends on the node’s opcode. See the Node docstring for more information.
Assignment to this property is allowed. All accounting of uses and users is updated automatically on assignment.


Returns the next Node in the linked list of Nodes.
The next Node in the linked list of Nodes.


Insert x before this node in the list of nodes in the graph. Example:
x (Node) – The node to put before this node. Must be a member of the same graph.


Returns the previous Node in the linked list of Nodes.
The previous Node in the linked list of Nodes.


Replace all uses of self in the Graph with the Node replace_with.
replace_with (Node) – The node to replace all uses of self with.
The list of Nodes on which this change was made.


Proxy objects are Node wrappers that flow through the program during symbolic tracing and record all the operations (torch function calls, method calls, operators) that they touch into the growing FX Graph.
If you’re doing graph transforms, you can wrap your own Proxy method around a raw Node so that you can use the overloaded operators to add additional things to a Graph.


Matches all possible non-overlapping sets of operators and their data dependencies (pattern) in the Graph of a GraphModule (gm), then replaces each of these matched subgraphs with another subgraph (replacement).
A list of Match objects representing the places in the original graph that pattern was matched to. The list is empty if there are no matches. Match is defined as:
List[Match]
Examples:
The above code will first match pattern in the forward method of traced_module. Pattern-matching is done based on use-def relationships, not node names. For example, if you had p = torch.cat([a, b]) in pattern, you could match m = torch.cat([a, b]) in the original forward function, despite the variable names being different (p vs m).
The return statement in pattern is matched based on its value only; it may or may not match to the return statement in the larger graph. In other words, the pattern doesn’t have to extend to the end of the larger graph.
When the pattern is matched, it will be removed from the larger function and replaced by replacement. If there are multiple matches for pattern in the larger function, each non-overlapping match will be replaced. In the case of a match overlap, the first found match in the set of overlapping matches will be replaced. (“First” here being defined as the first in a topological ordering of the Nodes’ use-def relationships. In most cases, the first Node is the parameter that appears directly after self, while the last Node is whatever the function returns.)
One important thing to note is that the parameters of the pattern Callable must be used in the Callable itself, and the parameters of the replacement Callable must match the pattern. The first rule is why, in the above code block, the forward function has parameters x, w1, w2, but the pattern function only has parameters w1, w2. pattern doesn’t use x, so it shouldn’t specify x as a parameter. As an example of the second rule, consider replacing
with
In this case, replacement needs the same number of parameters as pattern (both x and y), even though the parameter y isn’t used in replacement.
After calling subgraph_rewriter.replace_pattern, the generated Python code looks like this:


Symbolic tracing API
Given an nn.Module or function instance root, this function will return a GraphModule constructed by recording operations seen while tracing through root.
a Module created from the recorded operations from root.
GraphModule


Tracer is the class that implements the symbolic tracing functionality of torch.fx.symbolic_trace. A call to symbolic_trace(m) is equivalent to Tracer().trace(m).
Tracer can be subclassed to override various behaviors of the tracing process. The different behaviors that can be overridden are described in the docstrings of the methods on this class.
Method that specifies the behavior of this Tracer when it encounters a call to an nn.Module instance.
By default, the behavior is to check if the called module is a leaf module via is_leaf_module. If it is, emit a call_module node referring to m in the Graph. Otherwise, call the Module normally, tracing through the operations in its forward function.
This method can be overridden to–for example–create nested traced GraphModules, or any other behavior you would want while tracing across Module boundaries. Module boundaries.
The return value from the Module call. In the case that a call_module node was emitted, this is a Proxy value. Otherwise, it is whatever value was returned from the Module invocation.
A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the Graph.
By default, the behavior includes:
Given a non-Proxy Tensor object, emit IR for various cases:
This method can be overridden to support more types.
a (Any) – The value to be emitted as an Argument in the Graph.
The value a converted into the appropriate Argument
Create placeholder nodes corresponding to the signature of the root Module. This method introspects root’s signature and emits those nodes accordingly, also supporting *args and **kwargs.
A method to specify whether a given nn.Module is a “leaf” module.
Leaf modules are the atomic units that appear in the IR, referenced by call_module calls. By default, Modules in the PyTorch standard library namespace (torch.nn) are leaf modules. All other modules are traced through and their constituent ops are recorded, unless specified otherwise via this parameter.
Helper method to find the qualified name of mod in the Module hierarchy of root. For example, if root has a submodule named foo, which has a submodule named bar, passing bar into this function will return the string “foo.bar”.
mod (str) – The Module to retrieve the qualified name for.
Trace root and return the corresponding FX Graph representation. root can either be an nn.Module instance or a Python callable.
Note that after this call, self.root may be different from the root passed in here. For example, when a free function is passed to trace(), we will create an nn.Module instance to use as the root and add embedded constants to.
root (Union[Module, Callable]) – Either a Module or a function to be traced through.
A Graph representing the semantics of the passed-in root.


Method that specifies the behavior of this Tracer when it encounters a call to an nn.Module instance.
By default, the behavior is to check if the called module is a leaf module via is_leaf_module. If it is, emit a call_module node referring to m in the Graph. Otherwise, call the Module normally, tracing through the operations in its forward function.
This method can be overridden to–for example–create nested traced GraphModules, or any other behavior you would want while tracing across Module boundaries. Module boundaries.
The return value from the Module call. In the case that a call_module node was emitted, this is a Proxy value. Otherwise, it is whatever value was returned from the Module invocation.


A method to specify the behavior of tracing when preparing values to be used as arguments to nodes in the Graph.
By default, the behavior includes:
Given a non-Proxy Tensor object, emit IR for various cases:
This method can be overridden to support more types.
a (Any) – The value to be emitted as an Argument in the Graph.
The value a converted into the appropriate Argument


Create placeholder nodes corresponding to the signature of the root Module. This method introspects root’s signature and emits those nodes accordingly, also supporting *args and **kwargs.


A method to specify whether a given nn.Module is a “leaf” module.
Leaf modules are the atomic units that appear in the IR, referenced by call_module calls. By default, Modules in the PyTorch standard library namespace (torch.nn) are leaf modules. All other modules are traced through and their constituent ops are recorded, unless specified otherwise via this parameter.


Helper method to find the qualified name of mod in the Module hierarchy of root. For example, if root has a submodule named foo, which has a submodule named bar, passing bar into this function will return the string “foo.bar”.
mod (str) – The Module to retrieve the qualified name for.


Trace root and return the corresponding FX Graph representation. root can either be an nn.Module instance or a Python callable.
Note that after this call, self.root may be different from the root passed in here. For example, when a free function is passed to trace(), we will create an nn.Module instance to use as the root and add embedded constants to.
root (Union[Module, Callable]) – Either a Module or a function to be traced through.
A Graph representing the semantics of the passed-in root.


Transformer is a special type of interpreter that produces a new Module. It exposes a transform() method that returns the transformed Module. Transformer does not require arguments to run, as Interpreter does. Transformer works entirely symbolically.
Suppose we want to swap all instances of torch.neg with torch.sigmoid and vice versa (including their Tensor method equivalents). We could subclass Transformer like so:
module (GraphModule) – The Module to be transformed.
Execute a get_attr node. In Transformer, this is overridden to insert a new get_attr node into the output graph.
Execute a placeholder node. In Transformer, this is overridden to insert a new placeholder into the output graph.
Transform self.module and return the transformed GraphModule.


Execute a get_attr node. In Transformer, this is overridden to insert a new get_attr node into the output graph.


Execute a placeholder node. In Transformer, this is overridden to insert a new placeholder into the output graph.


Transform self.module and return the transformed GraphModule.


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 being traced through:
This function can also equivalently be used as a decorator:
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 trace rather than traced through.
fn_or_name (Union[str, Callable]) – The function or name of the global function to insert into the graph when it’s called


Gathers values along an axis specified by dim.
For a 3-D tensor the output is specified by:
input and index must have the same number of dimensions. It is also required that index.size(d) <= input.size(d) for all dimensions d != dim. out will have the same shape as index. Note that input and index do not broadcast against each other.
Example:


Computes the element-wise greatest common divisor (GCD) of input and other.
Both input and other must have integer types.
Note
This defines gcd(0,0)=0gcd(0, 0) = 0 .
out (Tensor, optional) – the output tensor.
Example:


Computes input≥other\text{input} \geq \text{other}  element-wise.
The second argument can be a number or a tensor whose shape is broadcastable with the first argument.
out (Tensor, optional) – the output tensor.
A boolean tensor that is True where input is greater than or equal to other and False elsewhere
Example:


Creates and returns a generator object that manages the state of the algorithm which produces pseudo random numbers. Used as a keyword argument in many In-place random sampling functions.
device (torch.device, optional) – the desired device for the generator.
An torch.Generator object.
Generator
Example:
Generator.device -> device
Gets the current device of the generator.
Example:
Returns the Generator state as a torch.ByteTensor.
A torch.ByteTensor which contains all the necessary bits to restore a Generator to a specific point in time.
Tensor
Example:
Returns the initial seed for generating random numbers.
Example:
Sets the seed for generating random numbers. Returns a torch.Generator object. It is recommended to set a large seed, i.e. a number that has a good balance of 0 and 1 bits. Avoid having many 0 bits in the seed.
seed (int) – The desired seed. Value must be within the inclusive range [-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]. Otherwise, a RuntimeError is raised. Negative inputs are remapped to positive values with the formula 0xffff_ffff_ffff_ffff + seed.
An torch.Generator object.
Generator
Example:
Gets a non-deterministic random number from std::random_device or the current time and uses it to seed a Generator.
Example:
Sets the Generator state.
new_state (torch.ByteTensor) – The desired state.
Example:


Generator.device -> device
Gets the current device of the generator.
Example:


Returns the Generator state as a torch.ByteTensor.
A torch.ByteTensor which contains all the necessary bits to restore a Generator to a specific point in time.
Tensor
Example:


Returns the initial seed for generating random numbers.
Example:


Sets the seed for generating random numbers. Returns a torch.Generator object. It is recommended to set a large seed, i.e. a number that has a good balance of 0 and 1 bits. Avoid having many 0 bits in the seed.
seed (int) – The desired seed. Value must be within the inclusive range [-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]. Otherwise, a RuntimeError is raised. Negative inputs are remapped to positive values with the formula 0xffff_ffff_ffff_ffff + seed.
An torch.Generator object.
Generator
Example:


Gets a non-deterministic random number from std::random_device or the current time and uses it to seed a Generator.
Example:


Sets the Generator state.
new_state (torch.ByteTensor) – The desired state.
Example:


This is a low-level function for calling LAPACK directly. This function returns a namedtuple (a, tau) as defined in LAPACK documentation for geqrf .
You’ll generally want to use torch.qr() instead.
Computes a QR decomposition of input, but without constructing QQ  and RR  as explicit separate matrices.
Rather, this directly calls the underlying LAPACK function ?geqrf which produces a sequence of ‘elementary reflectors’.
See LAPACK documentation for geqrf for further details.
input (Tensor) – the input matrix
out (tuple, optional) – the output tuple of (Tensor, Tensor)


Alias of torch.outer().
Warning
This function is deprecated and will be removed in a future PyTorch release. Use torch.outer() instead.


Get the current default floating point torch.dtype.
Example:


Returns the number of threads used for inter-op parallelism on CPU (e.g. in JIT interpreter)


Returns the number of threads used for parallelizing CPU operations


Returns the random number generator state as a torch.ByteTensor.


Computes input>other\text{input} > \text{other}  element-wise.
The second argument can be a number or a tensor whose shape is broadcastable with the first argument.
out (Tensor, optional) – the output tensor.
A boolean tensor that is True where input is greater than other and False elsewhere
Example:


Hamming window function.
where NN  is the full window size.
The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the NN  in above formula is in fact window_length+1\text{window\_length} + 1 . Also, we always have torch.hamming_window(L, periodic=True) equal to torch.hamming_window(L + 1, periodic=False)[:-1]).
Note
If window_length =1=1 , the returned window contains a single value 1.
Note
This is a generalized version of torch.hann_window().
A 1-D tensor of size (window_length,)(\text{window\_length},)  containing the window
Tensor


Hann window function.
where NN  is the full window size.
The input window_length is a positive integer controlling the returned window size. periodic flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like torch.stft(). Therefore, if periodic is true, the NN  in above formula is in fact window_length+1\text{window\_length} + 1 . Also, we always have torch.hann_window(L, periodic=True) equal to torch.hann_window(L + 1, periodic=False)[:-1]).
Note
If window_length =1=1 , the returned window contains a single value 1.
A 1-D tensor of size (window_length,)(\text{window\_length},)  containing the window
Tensor


Computes the Heaviside step function for each element in input. The Heaviside step function is defined as:
out (Tensor, optional) – the output tensor.
Example:


Computes the histogram of a tensor.
The elements are sorted into equal width bins between min and max. If min and max are both zero, the minimum and maximum values of the data are used.
Elements lower than min and higher than max are ignored.
out (Tensor, optional) – the output tensor.
Histogram represented as a tensor
Tensor
Example:


Performs a matrix multiplication of a sparse COO matrix mat1 and a strided matrix mat2. The result is a (1 + 1)-dimensional hybrid COO matrix.
{out} –


Stack tensors in sequence horizontally (column wise).
This is equivalent to concatenation along the first axis for 1-D tensors, and along the second axis for all other tensors.
tensors (sequence of Tensors) – sequence of tensors to concatenate
out (Tensor, optional) – the output tensor.
Example:


Pytorch Hub is a pre-trained model repository designed to facilitate research reproducibility.
Pytorch Hub supports publishing pre-trained models(model definitions and pre-trained weights) to a github repository by adding a simple hubconf.py file;
hubconf.py can have multiple entrypoints. Each entrypoint is defined as a python function (example: a pre-trained model you want to publish).
Here is a code snippet specifies an entrypoint for resnet18 model if we expand the implementation in pytorch/vision/hubconf.py. In most case importing the right function in hubconf.py is sufficient. Here we just want to use the expanded version as an example to show how it works. You can see the full script in pytorch/vision repo
Pytorch Hub provides convenient APIs to explore all available models in hub through torch.hub.list(), show docstring and examples through torch.hub.help() and load the pre-trained models using torch.hub.load().
List all entrypoints available in github hubconf.
a list of available entrypoint names
entrypoints
Show the docstring of entrypoint model.
Load a model from a github repo or a local directory.
Note: Loading a model is the typical use case, but this can also be used to for loading other objects such as tokenizers, loss functions, etc.
If source is 'github', repo_or_dir is expected to be of the form repo_owner/repo_name[:tag_name] with an optional tag/branch.
If source is 'local', repo_or_dir is expected to be a path to a local directory.
The output of the model callable when called with the given *args and **kwargs.
Download object at the given URL to a local path.
Loads the Torch serialized object at the given URL.
If downloaded file is a zip file, it will be automatically decompressed.
If the object is already present in model_dir, it’s deserialized and returned. The default value of model_dir is <hub_dir>/checkpoints where hub_dir is the directory returned by get_dir().
Note that *args and **kwargs in torch.hub.load() are used to instantiate a model. After you have loaded a model, how can you find out what you can do with the model? A suggested workflow is
To help users explore without referring to documentation back and forth, we strongly recommend repo owners make function help messages clear and succinct. It’s also helpful to include a minimal working example.
The locations are used in the order of
Get the Torch Hub cache directory used for storing downloaded models & weights.
If set_dir() is not called, default path is $TORCH_HOME/hub where environment variable $TORCH_HOME defaults to $XDG_CACHE_HOME/torch. $XDG_CACHE_HOME follows the X Design Group specification of the Linux filesystem layout, with a default value ~/.cache if the environment variable is not set.
Optionally set the Torch Hub directory used to save downloaded models & weights.
d (string) – path to a local folder to save downloaded models & weights.
By default, we don’t clean up files after loading it. Hub uses the cache by default if it already exists in the directory returned by get_dir().
Users can force a reload by calling hub.load(..., force_reload=True). This will delete the existing github folder and downloaded weights, reinitialize a fresh download. This is useful when updates are published to the same branch, users can keep up with the latest release.
Torch hub works by importing the package as if it was installed. There’re some side effects introduced by importing in Python. For example, you can see new items in Python caches sys.modules and sys.path_importer_cache which is normal Python behavior.
A known limitation that worth mentioning here is user CANNOT load two different branches of the same repo in the same python process. It’s just like installing two packages with the same name in Python, which is not good. Cache might join the party and give you surprises if you actually try that. Of course it’s totally fine to load them in separate processes.


Download object at the given URL to a local path.


Get the Torch Hub cache directory used for storing downloaded models & weights.
If set_dir() is not called, default path is $TORCH_HOME/hub where environment variable $TORCH_HOME defaults to $XDG_CACHE_HOME/torch. $XDG_CACHE_HOME follows the X Design Group specification of the Linux filesystem layout, with a default value ~/.cache if the environment variable is not set.


List all entrypoints available in github hubconf.
a list of available entrypoint names
entrypoints


Load a model from a github repo or a local directory.
Note: Loading a model is the typical use case, but this can also be used to for loading other objects such as tokenizers, loss functions, etc.
If source is 'github', repo_or_dir is expected to be of the form repo_owner/repo_name[:tag_name] with an optional tag/branch.
If source is 'local', repo_or_dir is expected to be a path to a local directory.
The output of the model callable when called with the given *args and **kwargs.


Loads the Torch serialized object at the given URL.
If downloaded file is a zip file, it will be automatically decompressed.
If the object is already present in model_dir, it’s deserialized and returned. The default value of model_dir is <hub_dir>/checkpoints where hub_dir is the directory returned by get_dir().


Optionally set the Torch Hub directory used to save downloaded models & weights.
d (string) – path to a local folder to save downloaded models & weights.


Given the legs of a right triangle, return its hypotenuse.
The shapes of input and other must be broadcastable.
out (Tensor, optional) – the output tensor.
Example:


Computes the zeroth order modified Bessel function of the first kind for each element of input.
input (Tensor) – the input tensor
out (Tensor, optional) – the output tensor.
Example:


Computes the regularized lower incomplete gamma function:
where both inputi\text{input}_i  and otheri\text{other}_i  are weakly positive and at least one is strictly positive. If both are zero or either is negative then outi=nan\text{out}_i=\text{nan} . Γ(⋅)\Gamma(\cdot)  in the equation above is the gamma function,
See torch.igammac() and torch.lgamma() for related functions.
Supports broadcasting to a common shape and float inputs.
Note
The backward pass with respect to input is not yet supported. Please open an issue on PyTorch’s Github to request it.
out (Tensor, optional) – the output tensor.
Example:


Computes the regularized upper incomplete gamma function:
where both inputi\text{input}_i  and otheri\text{other}_i  are weakly positive and at least one is strictly positive. If both are zero or either is negative then outi=nan\text{out}_i=\text{nan} . Γ(⋅)\Gamma(\cdot)  in the equation above is the gamma function,
See torch.igamma() and torch.lgamma() for related functions.
Supports broadcasting to a common shape and float inputs.
Note
The backward pass with respect to input is not yet supported. Please open an issue on PyTorch’s Github to request it.
out (Tensor, optional) – the output tensor.
Example:


Returns a new tensor containing imaginary values of the self tensor. The returned tensor and self share the same underlying storage.
Warning
imag() is only supported for tensors with complex dtypes.
input (Tensor) – the input tensor.


Returns a new tensor which indexes the input tensor along dimension dim using the entries in index which is a LongTensor.
The returned tensor has the same number of dimensions as the original tensor (input). The dimth dimension has the same size as the length of index; other dimensions have the same size as in the original tensor.
Note
The returned tensor does not use the same storage as the original tensor. If out has a different shape than expected, we silently change it to the correct shape, reallocating the underlying storage if necessary.
out (Tensor, optional) – the output tensor.
Example:


Returns the initial seed for generating random numbers as a Python long.


Computes the dot product for 1D tensors. For higher dimensions, sums the product of elements from input and other along their last dimension.
Note
If either input or other is a scalar, the result is equivalent to torch.mul(input, other).
If both input and other are non-scalars, the size of their last dimension must match and the result is equivalent to torch.tensordot(input, other, dims=([-1], [-1]))
out (Tensor, optional) – Optional output tensor to write result into. The output shape is input.shape[:-1] + other.shape[:-1].
Example:


Takes the inverse of the square matrix input. input can be batches of 2D square tensors, in which case this function would return a tensor composed of individual inverses.
Supports real and complex input.
Note
torch.inverse() is deprecated. Please use torch.linalg.inv() instead.
Note
Irrespective of the original strides, the returned tensors will be transposed, i.e. with strides like input.contiguous().transpose(-2, -1).stride()
input (Tensor) – the input tensor of size (∗,n,n)(*, n, n)  where * is zero or more batch dimensions
out (Tensor, optional) – the output tensor.
Examples:


Returns a new tensor with boolean elements representing if each element of input is “close” to the corresponding element of other. Closeness is defined as:
where input and other are finite. Where input and/or other are nonfinite they are close if and only if they are equal, with NaNs being considered equal to each other when equal_nan is True.
Examples:


Returns a new tensor with boolean elements representing if each element is finite or not.
Real values are finite when they are not NaN, negative infinity, or infinity. Complex values are finite when both their real and imaginary parts are finite.
input (Tensor): the input tensor.
A boolean tensor that is True where input is finite and False elsewhere
Example:


Tests if each element of input is infinite (positive or negative infinity) or not.
Note
Complex values are infinite when their real or imaginary part is infinite.
{input}
A boolean tensor that is True where input is infinite and False elsewhere
Example:


Returns a new tensor with boolean elements representing if each element of input is NaN or not. Complex values are considered NaN when either their real and/or imaginary part is NaN.
input (Tensor) – the input tensor.
A boolean tensor that is True where input is NaN and False elsewhere
Example:


Tests if each element of input is negative infinity or not.
input (Tensor) – the input tensor.
out (Tensor, optional) – the output tensor.


Tests if each element of input is positive infinity or not.
input (Tensor) – the input tensor.
out (Tensor, optional) – the output tensor.


Returns a new tensor with boolean elements representing if each element of input is real-valued or not. All real-valued types are considered real. Complex values are considered real when their imaginary part is 0.
input (Tensor) – the input tensor.
A boolean tensor that is True where input is real and False elsewhere
Example:


Inverse short time Fourier Transform. This is expected to be the inverse of stft(). It has the same parameters (+ additional optional parameter of length) and it should return the least squares estimation of the original signal. The algorithm will check using the NOLA condition ( nonzero overlap).
Important consideration in the parameters window and center so that the envelop created by the summation of all the windows is never zero at certain point in time. Specifically, ∑t=−∞∞∣w∣2[n−t×hop_length]=0\sum_{t=-\infty}^{\infty} |w|^2[n-t\times hop\_length] \cancel{=} 0 .
Since stft() discards elements at the end of the signal if they do not fit in a frame, istft may return a shorter signal than the original signal (can occur if center is False since the signal isn’t padded).
If center is True, then there will be padding e.g. 'constant', 'reflect', etc. Left padding can be trimmed off exactly because they can be calculated but right padding cannot be calculated without additional information.
Example: Suppose the last window is: [17, 18, 0, 0, 0] vs [18, 0, 0, 0, 0]
The n_fft, hop_length, win_length are all the same which prevents the calculation of right padding. These additional values could be zeros or a reflection of the signal so providing length could be useful. If length is None then padding will be aggressively removed (some loss of signal).
[1] D. W. Griffin and J. S. Lim, “Signal estimation from modified short-time Fourier transform,” IEEE Trans. ASSP, vol.32, no.2, pp.236-243, Apr. 1984.
input (Tensor) –
The input tensor. Expected to be output of stft(), can either be complex (channel, fft_size, n_frame), or real (channel, fft_size, n_frame, 2) where the channel dimension is optional.
Deprecated since version 1.8.0: Real input is deprecated, use complex inputs as returned by stft(..., return_complex=True) instead.
Least squares estimation of the original signal of size (…, signal_length)
Tensor


Returns True if the data type of input is a complex data type i.e., one of torch.complex64, and torch.complex128.
input (Tensor) – the input tensor.


Returns True if the data type of input is a floating point data type i.e., one of torch.float64, torch.float32, torch.float16, and torch.bfloat16.
input (Tensor) – the input tensor.


Returns True if the input is a single element tensor which is not equal to zero after type conversions. i.e. not equal to torch.tensor([0.]) or torch.tensor([0]) or torch.tensor([False]). Throws a RuntimeError if torch.numel() != 1 (even in case of sparse tensors).
input (Tensor) – the input tensor.
Examples:


Returns True if obj is a PyTorch storage object.
obj (Object) – Object to test


Returns True if obj is a PyTorch tensor.
Note that this function is simply doing isinstance(obj, Tensor). Using that isinstance check is better for typechecking with mypy, and more explicit - so it’s recommended to use that instead of is_tensor.
obj (Object) – Object to test


This decorator indicates that a method on an nn.Module is used as an entry point into a ScriptModule and should be compiled.
forward implicitly is assumed to be an entry point, so it does not need this decorator. Functions and methods called from forward are compiled as they are seen by the compiler, so they do not need this decorator either.
Example (using @torch.jit.export on a method):


Creates an asynchronous task executing func and a reference to the value of the result of this execution. fork will return immediately, so the return value of func may not have been computed yet. To force completion of the task and access the return value invoke torch.jit.wait on the Future. fork invoked with a func which returns T is typed as torch.jit.Future[T]. fork calls can be arbitrarily nested, and may be invoked with positional and keyword arguments. Asynchronous execution will only occur when run in TorchScript. If run in pure python, fork will not execute in parallel. fork will also not execute in parallel when invoked while tracing, however the fork and wait calls will be captured in the exported IR Graph. .. warning:
a reference to the execution of func. The value T can only be accessed by forcing completion of func through torch.jit.wait.
torch.jit.Future[T]
Example (fork a free function):
Example (fork a module method):


Freezing a ScriptModule will clone it and attempt to inline the cloned module’s submodules, parameters, and attributes as constants in the TorchScript IR Graph. By default, forward will be preserved, as well as attributes & methods specified in preserved_attrs. Additionally, any attribute that is modified within a preserved method will be preserved.
Freezing currently only accepts ScriptModules that are in eval mode.
Frozen ScriptModule.
Example (Freezing a simple module with a Parameter):
Example (Freezing a module with preserved attributes)
Note
If you’re not sure why an attribute is not being inlined as a constant, you can run dump_alias_db on frozen_module.forward.graph to see if freezing has detected the attribute is being modified.


This decorator indicates to the compiler that a function or method should be ignored and left as a Python function. This allows you to leave code in your model that is not yet TorchScript compatible. If called from TorchScript, ignored functions will dispatch the call to the Python interpreter. Models with ignored functions cannot be exported; use @torch.jit.unused instead.
Example (using @torch.jit.ignore on a method):
Example (using @torch.jit.ignore(drop=True) on a method):


This function provides for conatiner type refinement in TorchScript. It can refine parameterized containers of the List, Dict, Tuple, and Optional types. E.g. List[str], Dict[str, List[torch.Tensor]], Optional[Tuple[int,str,int]]. It can also refine basic types such as bools and ints that are available in TorchScript.
False otherwise with no new type refinement
bool
Example (using torch.jit.isinstance for type refinement): .. testcode:


Function that returns True when in compilation and False otherwise. This is useful especially with the @unused decorator to leave code in your model that is not yet TorchScript compatible. .. testcode:


Load a ScriptModule or ScriptFunction previously saved with torch.jit.save
All previously saved modules, no matter their device, are first loaded onto CPU, and then are moved to the devices they were saved from. If this fails (e.g. because the run time system doesn’t have certain devices), an exception is raised.
A ScriptModule object.
Example:


Save an offline version of this module for use in a separate process. The saved module serializes all of the methods, submodules, parameters, and attributes of this module. It can be loaded into the C++ API using torch::jit::load(filename) or into the Python API with torch.jit.load.
To be able to save a module, it must not make any calls to native Python functions. This means that all submodules must be subclasses of ScriptModule as well.
Danger
All modules, no matter their device, are always loaded onto the CPU during loading. This is different from torch.load()’s semantics and may change in the future.
Note
torch.jit.save attempts to preserve the behavior of some operators across versions. For example, dividing two integer tensors in PyTorch 1.5 performed floor division, and if the module containing that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6 its division behavior will be preserved. The same module saved in PyTorch 1.6 will fail to load in PyTorch 1.5, however, since the behavior of division changed in 1.6, and 1.5 does not know how to replicate the 1.6 behavior.
Example:


Scripting a function or nn.Module will inspect the source code, compile it as TorchScript code using the TorchScript compiler, and return a ScriptModule or ScriptFunction. TorchScript itself is a subset of the Python language, so not all features in Python work, but we provide enough functionality to compute on tensors and do control-dependent operations. For a complete guide, see the TorchScript Language Reference.
torch.jit.script can be used as a function for modules and functions, and as a decorator @torch.jit.script for TorchScript Classes and functions.
obj (callable, class, or nn.Module) – The nn.Module, function, or class type to compile.
If obj is nn.Module, script returns a ScriptModule object. The returned ScriptModule will have the same set of sub-modules and parameters as the original nn.Module. If obj is a standalone function, a ScriptFunction will be returned.
The @torch.jit.script decorator will construct a ScriptFunction by compiling the body of the function.
Example (scripting a function):
Scripting an nn.Module by default will compile the forward method and recursively compile any methods, submodules, and functions called by forward. If a nn.Module only uses features supported in TorchScript, no changes to the original module code should be necessary. script will construct ScriptModule that has copies of the attributes, parameters, and methods of the original module.
Example (scripting a simple module with a Parameter):
Example (scripting a module with traced submodules):
To compile a method other than forward (and recursively compile anything it calls), add the @torch.jit.export decorator to the method. To opt out of compilation use @torch.jit.ignore or @torch.jit.unused.
Example (an exported and ignored method in a module):


Functionally equivalent to a ScriptModule, but represents a single function and does not have any attributes or Parameters.