diff --git a/.jenkins/validate_tutorials_built.py b/.jenkins/validate_tutorials_built.py index 75dd51dd78..c873c43d4d 100644 --- a/.jenkins/validate_tutorials_built.py +++ b/.jenkins/validate_tutorials_built.py @@ -23,7 +23,6 @@ "intermediate_source/dqn_with_rnn_tutorial", #not working on 2.8 release reenable after 3514 "intermediate_source/mnist_train_nas", # used by ax_multiobjective_nas_tutorial.py "intermediate_source/torch_compile_conv_bn_fuser", - "intermediate_source/_torch_export_nightly_tutorial", # does not work on release "advanced_source/usb_semisup_learn", # fails with CUDA OOM error, should try on a different worker "unstable_source/gpu_direct_storage", # requires specific filesystem + GPUDirect Storage to be set up "recipes_source/recipes/tensorboard_with_pytorch", diff --git a/conf.py b/conf.py index 9bfebed72b..8a760a7cdc 100644 --- a/conf.py +++ b/conf.py @@ -181,7 +181,6 @@ def wrapper(*args, **kwargs): "# https://docs.pytorch.org/tutorials/beginner/colab\n" "%matplotlib inline" ), - "ignore_pattern": r"_torch_export_nightly_tutorial.py", "pypandoc": { "extra_args": ["--mathjax", "--toc"], "filters": [".jenkins/custom_pandoc_filter.py"], diff --git a/intermediate_source/_torch_export_nightly_tutorial.py b/intermediate_source/_torch_export_nightly_tutorial.py deleted file mode 100644 index fdbe18392e..0000000000 --- a/intermediate_source/_torch_export_nightly_tutorial.py +++ /dev/null @@ -1,635 +0,0 @@ -# -*- coding: utf-8 -*- - -""" -torch.export Nightly Tutorial -================ -**Author:** William Wen, Zhengxu Chen, Angela Yi -""" - -###################################################################### -# -# .. warning:: -# -# ``torch.export`` and its related features are in prototype status and are subject to backwards compatibility -# breaking changes. This tutorial provides a snapshot of ``torch.export`` usage as of PyTorch 2.1. -# -# :func:`torch.export` is the PyTorch 2.X way to export PyTorch models into -# standardized model representations, intended -# to be run on different (i.e. Python-less) environments. -# -# In this tutorial, you will learn how to use :func:`torch.export` to extract -# ``ExportedProgram``'s (i.e. single-graph representations) from PyTorch programs. -# We also detail some considerations/modifications that you may need -# to make in order to make your model compatible with ``torch.export``. -# -# **Contents** -# -# .. contents:: -# :local: - -###################################################################### -# Basic Usage -# ----------- -# -# ``torch.export`` extracts single-graph representations from PyTorch programs -# by tracing the target function, given example inputs. -# ``torch.export.export()`` is the main entry point for ``torch.export``. -# -# In this tutorial, ``torch.export`` and ``torch.export.export()`` are practically synonymous, -# though ``torch.export`` generally refers to the PyTorch 2.X export process, and ``torch.export.export()`` -# generally refers to the actual function call. -# -# The signature of ``torch.export.export()`` is: -# -# .. code:: python -# -# export( -# f: Callable, -# args: Tuple[Any, ...], -# kwargs: Optional[Dict[str, Any]] = None, -# *, -# dynamic_shapes: Optional[Dict[str, Dict[int, Dim]]] = None -# ) -> ExportedProgram -# -# ``torch.export.export()`` traces the tensor computation graph from calling ``f(*args, **kwargs)`` -# and wraps it in an ``ExportedProgram``, which can be serialized or executed later with -# different inputs. Note that while the output ``ExportedGraph`` is callable and can be -# called in the same way as the original input callable, it is not a ``torch.nn.Module``. -# We will detail the ``dynamic_shapes`` argument later in the tutorial. - -import torch -from torch.export import export - -class MyModule(torch.nn.Module): - def __init__(self): - super().__init__() - self.lin = torch.nn.Linear(100, 10) - - def forward(self, x, y): - return torch.nn.functional.relu(self.lin(x + y), inplace=True) - -mod = MyModule() -exported_mod = export(mod, (torch.randn(8, 100), torch.randn(8, 100))) -print(type(exported_mod)) -print(exported_mod(torch.randn(8, 100), torch.randn(8, 100))) - -###################################################################### -# Let's review some attributes of ``ExportedProgram`` that are of interest. -# -# The ``graph`` attribute is an `FX graph `__ -# traced from the function we exported, that is, the computation graph of all PyTorch operations. -# The FX graph has some important properties: -# -# - The operations are "ATen-level" operations. -# - The graph is "functionalized", meaning that no operations are mutations. -# -# The ``graph_module`` attribute is the ``GraphModule`` that wraps the ``graph`` attribute -# so that it can be ran as a ``torch.nn.Module``. - -print(exported_mod) -print(exported_mod.graph_module) - -###################################################################### -# The printed code shows that FX graph only contains ATen-level ops (such as ``torch.ops.aten``) -# and that mutations were removed. For example, the mutating op ``torch.nn.functional.relu(..., inplace=True)`` -# is represented in the printed code by ``torch.ops.aten.relu.default``, which does not mutate. -# Future uses of input to the original mutating ``relu`` op are replaced by the additional new output -# of the replacement non-mutating ``relu`` op. -# -# Other attributes of interest in ``ExportedProgram`` include: -# -# - ``graph_signature`` -- the inputs, outputs, parameters, buffers, etc. of the exported graph. -# - ``range_constraints`` and ``equality_constraints`` -- constraints, covered later - -print(exported_mod.graph_signature) - -###################################################################### -# See the ``torch.export`` `documentation `__ -# for more details. - -###################################################################### -# Graph Breaks -# ------------ -# -# Although ``torch.export`` shares components with ``torch.compile``, -# the key limitation of ``torch.export``, especially when compared to ``torch.compile``, is that it does not -# support graph breaks. This is because handling graph breaks involves interpreting -# the unsupported operation with default Python evaluation, which is incompatible -# with the export use case. Therefore, in order to make your model code compatible -# with ``torch.export``, you will need to modify your code to remove graph breaks. -# -# A graph break is necessary in cases such as: -# -# - data-dependent control flow - -def bad1(x): - if x.sum() > 0: - return torch.sin(x) - return torch.cos(x) - -import traceback as tb -try: - export(bad1, (torch.randn(3, 3),)) -except Exception: - tb.print_exc() - -###################################################################### -# - accessing tensor data with ``.data`` - -def bad2(x): - x.data[0, 0] = 3 - return x - -try: - export(bad2, (torch.randn(3, 3),)) -except Exception: - tb.print_exc() - -###################################################################### -# - calling unsupported functions (such as many built-in functions) - -def bad3(x): - x = x + 1 - return x + id(x) - -try: - export(bad3, (torch.randn(3, 3),)) -except Exception: - tb.print_exc() - -###################################################################### -# - unsupported Python language features (e.g. throwing exceptions, match statements) - -def bad4(x): - try: - x = x + 1 - raise RuntimeError("bad") - except: - x = x + 2 - return x - -try: - export(bad4, (torch.randn(3, 3),)) -except Exception: - tb.print_exc() - -###################################################################### -# The sections below demonstrate some ways you can modify your code -# in order to remove graph breaks. - -###################################################################### -# Control Flow Ops -# ---------------- -# -# ``torch.export`` actually does support data-dependent control flow. -# But these need to be expressed using control flow ops. For example, -# we can fix the control flow example above using the ``cond`` op, like so: - -from functorch.experimental.control_flow import cond - -def bad1_fixed(x): - def true_fn(x): - return torch.sin(x) - def false_fn(x): - return torch.cos(x) - return cond(x.sum() > 0, true_fn, false_fn, [x]) - -exported_bad1_fixed = export(bad1_fixed, (torch.randn(3, 3),)) -print(exported_bad1_fixed(torch.ones(3, 3))) -print(exported_bad1_fixed(-torch.ones(3, 3))) - -###################################################################### -# There are limitations to ``cond`` that one should be aware of: -# -# - The predicate (i.e. ``x.sum() > 0``) must result in a boolean or a single-element tensor. -# - The operands (i.e. ``[x]``) must be tensors. -# - The branch function (i.e. ``true_fn`` and ``false_fn``) signature must match with the -# operands and they must both return a single tensor with the same metadata (for example, ``dtype``, ``shape``, etc.). -# - Branch functions cannot mutate input or global variables. -# - Branch functions cannot access closure variables, except for ``self`` if the function is -# defined in the scope of a method. -# -# For more details about ``cond``, check out the `documentation `__. - -###################################################################### -# .. -# [NOTE] map is not documented at the moment -# We can also use ``map``, which applies a function across the first dimension -# of the first tensor argument. -# -# from functorch.experimental.control_flow import map -# -# def map_example(xs): -# def map_fn(x, const): -# def true_fn(x): -# return x + const -# def false_fn(x): -# return x - const -# return control_flow.cond(x.sum() > 0, true_fn, false_fn, [x]) -# return control_flow.map(map_fn, xs, torch.tensor([2.0])) -# -# exported_map_example= export(map_example, (torch.randn(4, 3),)) -# inp = torch.cat((torch.ones(2, 3), -torch.ones(2, 3))) -# print(exported_map_example(inp)) - -###################################################################### -# Constraints/Dynamic Shapes -# -------------------------- -# -# Ops can have different specializations/behaviors for different tensor shapes, so by default, -# ``torch.export`` requires inputs to ``ExportedProgram`` to have the same shape as the respective -# example inputs given to the initial ``torch.export.export()`` call. -# If we try to run the ``ExportedProgram`` in the example below with a tensor -# with a different shape, we get an error: - -class MyModule2(torch.nn.Module): - def __init__(self): - super().__init__() - self.lin = torch.nn.Linear(100, 10) - - def forward(self, x, y): - return torch.nn.functional.relu(self.lin(x + y), inplace=True) - -mod2 = MyModule2() -exported_mod2 = export(mod2, (torch.randn(8, 100), torch.randn(8, 100))) - -try: - exported_mod2(torch.randn(10, 100), torch.randn(10, 100)) -except Exception: - tb.print_exc() - -###################################################################### -# We can relax this constraint using the ``dynamic_shapes`` argument of -# ``torch.export.export()``, which allows us to specify, using ``torch.export.Dim`` -# (`documentation `__), -# which dimensions of the input tensors are dynamic. -# -# For each tensor argument of the input callable, we can specify a mapping from the dimension -# to a ``torch.export.Dim``. -# A ``torch.export.Dim`` is essentially a named symbolic integer with optional -# minimum and maximum bounds. -# -# Then, the format of ``torch.export.export()``'s ``dynamic_shapes`` argument is a mapping -# from the input callable's tensor argument names, to dimension --> dim mappings as described above. -# If there is no ``torch.export.Dim`` given to a tensor argument's dimension, then that dimension is -# assumed to be static. -# -# The first argument of ``torch.export.Dim`` is the name for the symbolic integer, used for debugging. -# Then we can specify an optional minimum and maximum bound (inclusive). Below, we show example usage. -# -# In the example below, our input -# ``inp1`` has an unconstrained first dimension, but the size of the second -# dimension must be in the interval [4, 18]. - -from torch.export import Dim - -inp1 = torch.randn(10, 10, 2) - -def dynamic_shapes_example1(x): - x = x[:, 2:] - return torch.relu(x) - -inp1_dim0 = Dim("inp1_dim0") -inp1_dim1 = Dim("inp1_dim1", min=4, max=18) -dynamic_shapes1 = { - "x": {0: inp1_dim0, 1: inp1_dim1}, -} - -exported_dynamic_shapes_example1 = export(dynamic_shapes_example1, (inp1,), dynamic_shapes=dynamic_shapes1) - -print(exported_dynamic_shapes_example1(torch.randn(5, 5, 2))) - -try: - exported_dynamic_shapes_example1(torch.randn(8, 1, 2)) -except Exception: - tb.print_exc() - -try: - exported_dynamic_shapes_example1(torch.randn(8, 20, 2)) -except Exception: - tb.print_exc() - -try: - exported_dynamic_shapes_example1(torch.randn(8, 8, 3)) -except Exception: - tb.print_exc() - -###################################################################### -# Note that if our example inputs to ``torch.export`` do not satisfy the constraints -# given by ``dynamic_shapes``, then we get an error. - -inp1_dim1_bad = Dim("inp1_dim1_bad", min=11, max=18) -dynamic_shapes1_bad = { - "x": {0: inp1_dim0, 1: inp1_dim1_bad}, -} - -try: - export(dynamic_shapes_example1, (inp1,), dynamic_shapes=dynamic_shapes1_bad) -except Exception: - tb.print_exc() - -###################################################################### -# We can enforce that equalities between dimensions of different tensors -# by using the same ``torch.export.Dim`` object, for example, in matrix multiplication: - -inp2 = torch.randn(4, 8) -inp3 = torch.randn(8, 2) - -def dynamic_shapes_example2(x, y): - return x @ y - -inp2_dim0 = Dim("inp2_dim0") -inner_dim = Dim("inner_dim") -inp3_dim1 = Dim("inp3_dim1") - -dynamic_shapes2 = { - "x": {0: inp2_dim0, 1: inner_dim}, - "y": {0: inner_dim, 1: inp3_dim1}, -} - -exported_dynamic_shapes_example2 = export(dynamic_shapes_example2, (inp2, inp3), dynamic_shapes=dynamic_shapes2) - -print(exported_dynamic_shapes_example2(torch.randn(2, 16), torch.randn(16, 4))) - -try: - exported_dynamic_shapes_example2(torch.randn(4, 8), torch.randn(4, 2)) -except Exception: - tb.print_exc() - -###################################################################### -# We can actually use ``torch.export`` to guide us as to which ``dynamic_shapes`` constraints -# are necessary. We can do this by relaxing all constraints (recall that if we -# do not provide constraints for a dimension, the default behavior is to constrain -# to the exact shape value of the example input) and letting ``torch.export`` -# error out. - -inp4 = torch.randn(8, 16) -inp5 = torch.randn(16, 32) - -def dynamic_shapes_example3(x, y): - if x.shape[0] <= 16: - return x @ y[:, :16] - return y - -dynamic_shapes3 = { - "x": {i: Dim(f"inp4_dim{i}") for i in range(inp4.dim())}, - "y": {i: Dim(f"inp5_dim{i}") for i in range(inp5.dim())}, -} - -try: - export(dynamic_shapes_example3, (inp4, inp5), dynamic_shapes=dynamic_shapes3) -except Exception: - tb.print_exc() - -###################################################################### -# We can see that the error message gives us suggested fixes to our -# dynamic shape constraints. Let us follow those suggestions (exact -# suggestions may differ slightly): - -def suggested_fixes(): - inp4_dim1 = Dim('shared_dim') - # suggested fixes below - inp4_dim0 = Dim('inp4_dim0', max=16) - inp5_dim1 = Dim('inp5_dim1', min=17) - inp5_dim0 = inp4_dim1 - # end of suggested fixes - return { - "x": {0: inp4_dim0, 1: inp4_dim1}, - "y": {0: inp5_dim0, 1: inp5_dim1}, - } - -dynamic_shapes3_fixed = suggested_fixes() -exported_dynamic_shapes_example3 = export(dynamic_shapes_example3, (inp4, inp5), dynamic_shapes=dynamic_shapes3_fixed) -print(exported_dynamic_shapes_example3(torch.randn(4, 32), torch.randn(32, 64))) - -###################################################################### -# Note that in the example above, because we constrained the value of ``x.shape[0]`` in -# ``dynamic_shapes_example3``, the exported program is sound even though there is a -# raw ``if`` statement. -# -# If you want to see why ``torch.export`` generated these constraints, you can -# re-run the script with the environment variable ``TORCH_LOGS=dynamic,dynamo``, -# or use ``torch._logging.set_logs``. - -import logging -torch._logging.set_logs(dynamic=logging.INFO, dynamo=logging.INFO) -exported_dynamic_shapes_example3 = export(dynamic_shapes_example3, (inp4, inp5), dynamic_shapes=dynamic_shapes3_fixed) - -# reset to previous values -torch._logging.set_logs(dynamic=logging.WARNING, dynamo=logging.WARNING) - -###################################################################### -# We can view an ``ExportedProgram``'s constraints using the ``range_constraints`` and -# ``equality_constraints`` attributes. The logging above reveals what the symbols ``s0, s1, ...`` -# represent. - -print(exported_dynamic_shapes_example3.range_constraints) -print(exported_dynamic_shapes_example3.equality_constraints) - -###################################################################### -# Custom Ops -# ---------- -# -# ``torch.export`` can export PyTorch programs with custom operators. -# -# -# Currently, the steps to register a custom op for use by ``torch.export`` are: -# -# - If you’re writing custom ops purely in Python, use torch.library.custom_op. - -import torch.library -import numpy as np - -@torch.library.custom_op("mylib::sin", mutates_args=()) -def sin(x): - x_np = x.numpy() - y_np = np.sin(x_np) - return torch.from_numpy(y_np) - -###################################################################### -# - You will need to provide abstract implementation so that PT2 can trace through it. - -@torch.library.register_fake("mylib::sin") -def _(x): - return torch.empty_like(x) - -# - Sometimes, the custom op you are exporting has data-dependent output, meaning -# we can't determine the shape of the output at compile time. In this case, you can do -# following: -@torch.library.custom_op("mylib::nonzero", mutates_args=()) -def nonzero(x): - x_np = x.cpu().numpy() - res = np.stack(np.nonzero(x_np), axis=1) - return torch.tensor(res, device=x.device) - -@torch.library.register_fake("mylib::nonzero") -def _(x): - # The number of nonzero-elements is data-dependent. - # Since we cannot peek at the data in an abstract implementation, - # we use the `ctx` object to construct a new ``symint`` that - # represents the data-dependent size. - ctx = torch.library.get_ctx() - nnz = ctx.new_dynamic_size() - shape = [nnz, x.dim()] - result = x.new_empty(shape, dtype=torch.int64) - return result - -###################################################################### -# - Call the custom op from the code you want to export using ``torch.ops`` - -def custom_op_example(x): - x = torch.sin(x) - x = torch.ops.mylib.sin(x) - x = torch.cos(x) - y = torch.ops.mylib.nonzero(x) - return x + y.sum() - -###################################################################### -# - Export the code as before - -exported_custom_op_example = export(custom_op_example, (torch.randn(3, 3),)) -exported_custom_op_example.graph_module.print_readable() -print(exported_custom_op_example(torch.randn(3, 3))) - -###################################################################### -# Note in the above outputs that the custom op is included in the exported graph. -# And when we call the exported graph as a function, the original custom op is called, -# as evidenced by the ``print`` call. -# -# If you have a custom operator implemented in C++, please refer to -# `this document `__ -# to make it compatible with ``torch.export``. - -###################################################################### -# Decompositions -# -------------- -# -# The graph produced by ``torch.export`` by default returns a graph containing -# only functional ATen operators. This functional ATen operator set (or "opset") contains around 2000 -# operators, all of which are functional, that is, they do not -# mutate or alias inputs. You can find a list of all ATen operators -# `here `__ -# and you can inspect if an operator is functional by checking -# ``op._schema.is_mutable``, for example: - -print(torch.ops.aten.add.Tensor._schema.is_mutable) -print(torch.ops.aten.add_.Tensor._schema.is_mutable) - -###################################################################### -# By default, the environment in which you want to run the exported graph -# should support all ~2000 of these operators. -# However, you can use the following API on the exported program -# if your specific environment is only able to support a subset of -# the ~2000 operators. -# -# .. code:: python -# -# def run_decompositions( -# self: ExportedProgram, -# decomposition_table: Optional[Dict[torch._ops.OperatorBase, Callable]] -# ) -> ExportedProgram -# -# ``run_decompositions`` takes in a decomposition table, which is a mapping of -# operators to a function specifying how to reduce, or decompose, that operator -# into an equivalent sequence of other ATen operators. -# -# The default decomposition table for ``run_decompositions`` is the -# `Core ATen decomposition table `__ -# which will decompose the all ATen operators to the -# `Core ATen Operator Set `__ -# which consists of only ~180 operators. - -class M(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(3, 4) - - def forward(self, x): - return self.linear(x) - -ep = export(M(), (torch.randn(2, 3),)) -print(ep.graph) - -core_ir_ep = ep.run_decompositions() -print(core_ir_ep.graph) - -###################################################################### -# Notice that after running ``run_decompositions`` the -# ``torch.ops.aten.t.default`` operator, which is not part of the Core ATen -# Opset, has been replaced with ``torch.ops.aten.permute.default`` which is part -# of the Core ATen Opset. - -###################################################################### -# Most ATen operators already have decompositions, which are located -# `here `__. -# If you would like to use some of these existing decomposition functions, -# you can pass in a list of operators you would like to decompose to the -# `get_decompositions `__ -# function, which will return a decomposition table using existing -# decomposition implementations. - -class M(torch.nn.Module): - def __init__(self): - super().__init__() - self.linear = torch.nn.Linear(3, 4) - - def forward(self, x): - return self.linear(x) - -ep = export(M(), (torch.randn(2, 3),)) -print(ep.graph) - -from torch._decomp import get_decompositions -decomp_table = get_decompositions([torch.ops.aten.t.default, torch.ops.aten.transpose.int]) -core_ir_ep = ep.run_decompositions(decomp_table) -print(core_ir_ep.graph) - -###################################################################### -# If there is no existing decomposition function for an ATen operator that you would -# like to decompose, feel free to send a pull request into PyTorch -# implementing the decomposition! - -###################################################################### -# ExportDB -# -------- -# -# ``torch.export`` will only ever export a single computation graph from a PyTorch program. Because of this requirement, -# there will be Python or PyTorch features that are not compatible with ``torch.export``, which will require users to -# rewrite parts of their model code. We have seen examples of this earlier in the tutorial -- for example, rewriting -# if-statements using ``cond``. -# -# `ExportDB `__ is the standard reference that documents -# supported and unsupported Python/PyTorch features for ``torch.export``. It is essentially a list a program samples, each -# of which represents the usage of one particular Python/PyTorch feature and its interaction with ``torch.export``. -# Examples are also tagged by category so that they can be more easily searched. -# -# For example, let's use ExportDB to get a better understanding of how the predicate works in the ``cond`` operator. -# We can look at the example called ``cond_predicate``, which has a ``torch.cond`` tag. The example code looks like: - -def cond_predicate(x): - """ - The conditional statement (aka predicate) passed to ``cond()`` must be one of the following: - - torch.Tensor with a single element - - boolean expression - NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized. - """ - pred = x.dim() > 2 and x.shape[2] > 10 - return cond(pred, lambda x: x.cos(), lambda y: y.sin(), [x]) - -###################################################################### -# More generally, ExportDB can be used as a reference when one of the following occurs: -# -# 1. Before attempting ``torch.export``, you know ahead of time that your model uses some tricky Python/PyTorch features -# and you want to know if ``torch.export`` covers that feature. -# 2. When attempting ``torch.export``, there is a failure and it's unclear how to work around it. -# -# ExportDB is not exhaustive, but is intended to cover all use cases found in typical PyTorch code. Feel free to reach -# out if there is an important Python/PyTorch feature that should be added to ExportDB or supported by ``torch.export``. - -###################################################################### -# Conclusion -# ---------- -# -# We introduced ``torch.export``, the new PyTorch 2.X way to export single computation -# graphs from PyTorch programs. In particular, we demonstrate several code modifications -# and considerations (control flow ops, constraints, etc.) that need to be made in order to export a graph. diff --git a/intermediate_source/torch_export_nightly_tutorial.rst b/intermediate_source/torch_export_nightly_tutorial.rst deleted file mode 100644 index e7ef2e8815..0000000000 --- a/intermediate_source/torch_export_nightly_tutorial.rst +++ /dev/null @@ -1,10 +0,0 @@ -torch.export Nightly Tutorial -============================= - -This tutorial has been moved to https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html - -It will redirect in 3 seconds. - -.. raw:: html - - diff --git a/intermediate_source/torch_export_tutorial.py b/intermediate_source/torch_export_tutorial.py index 3ca6d09a52..3c2fa67bb9 100644 --- a/intermediate_source/torch_export_tutorial.py +++ b/intermediate_source/torch_export_tutorial.py @@ -16,7 +16,7 @@ # :func:`torch.export` is the PyTorch 2.X way to export PyTorch models into # standardized model representations, intended # to be run on different (i.e. Python-less) environments. The official -# documentation can be found `here `__. +# documentation can be found `here `__. # # In this tutorial, you will learn how to use :func:`torch.export` to extract # ``ExportedProgram``'s (i.e. single-graph representations) from PyTorch programs. @@ -79,7 +79,7 @@ def forward(self, x, y): ###################################################################### # Let's review some attributes of ``ExportedProgram`` that are of interest. # -# The ``graph`` attribute is an `FX graph `__ +# The ``graph`` attribute is an `FX graph `__ # traced from the function we exported, that is, the computation graph of all PyTorch operations. # The FX graph is in "ATen IR" meaning that it contains only "ATen-level" operations. # @@ -92,7 +92,7 @@ def forward(self, x, y): print(exported_mod) ###################################################################### -# See the ``torch.export`` `documentation `__ +# See the ``torch.export`` `documentation `__ # for more details. ###################################################################### @@ -220,7 +220,7 @@ def false_fn(x): # - Branch functions cannot access closure variables, except for ``self`` if the function is # defined in the scope of a method. # -# For more details about ``cond``, check out the `cond documentation `__. +# For more details about ``cond``, check out the `cond documentation `__. ###################################################################### # We can also use ``map``, which applies a function across the first dimension @@ -308,7 +308,7 @@ def forward( ###################################################################### # Before we look at the program that's produced, let's understand what specifying ``dynamic_shapes`` entails, # and how that interacts with export. For every input dimension where a ``Dim`` object is specified, a symbol is -# `allocated `_, +# `allocated `_, # taking on a range of ``[2, inf]`` (why not ``[0, inf]`` or ``[1, inf]``? we'll explain later in the # 0/1 specialization section). # @@ -605,7 +605,7 @@ def forward(self, x, y): # How are these values represented in the exported program? In the `Constraints/Dynamic Shapes `_ # section, we talked about allocating symbols to represent dynamic input dimensions. # The same happens here: we allocate symbols for every data-dependent value that appears in the program. The important distinction is that these are "unbacked" symbols, -# in contrast to the "backed" symbols allocated for input dimensions. The `"backed/unbacked" `_ +# in contrast to the "backed" symbols allocated for input dimensions. The `"backed/unbacked" `_ # nomenclature refers to the presence/absence of a "hint" for the symbol: a concrete value backing the symbol, that can inform the compiler on how to proceed. # # In the input shape symbol case (backed symbols), these hints are simply the sample input shapes provided, which explains why control-flow branching is determined by the sample input properties. @@ -637,7 +637,7 @@ def forward(self, x, y): # ^^^^^^^^^^^^^^^^^^^^^^ # # But the case above is easy to export, because the concrete values of these symbols aren't used in any compiler decision-making; all that's relevant is that the return values are unbacked symbols. -# The data-dependent errors highlighted in this section are cases like the following, where `data-dependent guards `_ are encountered: +# The data-dependent errors highlighted in this section are cases like the following, where `data-dependent guards `_ are encountered: class Foo(torch.nn.Module): def forward(self, x, y): @@ -779,7 +779,7 @@ def forward(self, x, y): ###################################################################### # Data-dependent errors can be much more involved, and there are many more options in your toolkit to deal with them: ``torch._check_is_size()``, ``guard_size_oblivious()``, or real-tensor tracing, as starters. -# For more in-depth guides, please refer to the `Export Programming Model `_, +# For more in-depth guides, please refer to the `Export Programming Model `_, # or `Dealing with GuardOnDataDependentSymNode errors `_. ###################################################################### @@ -787,7 +787,7 @@ def forward(self, x, y): # ---------- # # ``torch.export`` can export PyTorch programs with custom operators. Please -# refer to `this page `__ +# refer to `this page `__ # on how to author a custom operator in either C++ or Python. # # The following is an example of registering a custom operator in python to be @@ -843,10 +843,7 @@ def forward(self, x): print(torch.ops.aten.add_.Tensor._schema.is_mutable) ###################################################################### -# This generic IR can be used to train in eager PyTorch Autograd. This IR can be -# more explicitly reached through the API ``torch.export.export_for_training``, -# which was introduced in PyTorch 2.5, but calling ``torch.export.export`` -# should produce the same graph as of PyTorch 2.6. +# This generic IR can be used to train in eager PyTorch Autograd. class DecompExample(torch.nn.Module): def __init__(self) -> None: @@ -859,7 +856,7 @@ def forward(self, x): x = self.bn(x) return (x,) -ep_for_training = torch.export.export_for_training(DecompExample(), (torch.randn(1, 1, 3, 3),)) +ep_for_training = torch.export.export(DecompExample(), (torch.randn(1, 1, 3, 3),)) print(ep_for_training.graph) ###################################################################### @@ -882,7 +879,7 @@ def forward(self, x): ###################################################################### # We can also further lower this exported program to an operator set which only # contains the -# `Core ATen Operator Set `__, +# `Core ATen Operator Set `__, # which is a collection of only ~180 operators. This IR is optimal for backends # who do not want to reimplement all ATen operators. @@ -925,7 +922,7 @@ def my_awesome_custom_conv2d_function(x, weight, bias, stride=[1, 1], padding=[0 # rewrite parts of their model code. We have seen examples of this earlier in the tutorial -- for example, rewriting # if-statements using ``cond``. # -# `ExportDB `__ is the standard reference that documents +# `ExportDB `__ is the standard reference that documents # supported and unsupported Python/PyTorch features for ``torch.export``. It is essentially a list a program samples, each # of which represents the usage of one particular Python/PyTorch feature and its interaction with ``torch.export``. # Examples are also tagged by category so that they can be more easily searched. @@ -961,7 +958,7 @@ def cond_predicate(x): # produced by ``torch.export`` eagerly will be equivalent to running the eager # module. To optimize the execution of the Exported Program, we can pass this # exported artifact to backends such as Inductor through ``torch.compile``, -# `AOTInductor `__, +# `AOTInductor `__, # or `TensorRT `__. class M(torch.nn.Module): @@ -997,7 +994,7 @@ def forward(self, x): # # # Load and run the .so file in Python. # # To load and run it in a C++ environment, see: -# # https://pytorch.org/docs/main/torch.compiler_aot_inductor.html +# # https://docs.pytorch.org/docs/main/torch.compiler_aot_inductor.html # aoti_compiled = torch._inductor.aoti_load_package(pt2_path) # res = aoti_compiled(inp)