[export] Update docs (pytorch#157750)

angelayi · pytorchmergebot · commit 1cc62c2cb91e · 2025-07-16T19:53:12.000Z
Preview: https://docs-preview.pytorch.org/pytorch/pytorch/157750/export.html Changes: * Rename draft_export.md -> export.draft_export.md for consistency. * Removed non-strict section in export, instead pointed to programming model doc. * Extended "Expressing Dynamism" section to include Dim hints, ShapeCollection, and AdditionalInputs. * Removed Specialization section in favor of programming model doc * Added pt2 archive doc * Cleaned up sidebar Pull Request resolved: pytorch#157750 Approved by: https://github.com/pianpwk
diff --git a/docs/source/export.md b/docs/source/export.md
diff --git a/docs/source/export/draft_export.md b/docs/source/export/draft_export.md
@@ -1,4 +1,4 @@
-(draft-export)=
+(export.draft_export)=
 
 # Draft Export
 
@@ -126,7 +126,7 @@ Running the `tlparse` command in the terminal will generate a
 [tlparse](https://github.com/pytorch/tlparse)
 HTML report. Here is an example of the `tlparse` report:
 
-```{image} _static/img/export/draft_export_report.png
+```{image} ../_static/img/export/draft_export_report.png
 ```
 
 Clicking into the Data Dependent Error, we will see the following page which
@@ -136,7 +136,7 @@ contains information to help debug this error. Specifically, it contains:
 - A list of local variables and their shapes
 - Information for how this guard was created
 
-```{image} _static/img/export/draft_export_report_dde.png
+```{image} ../_static/img/export/draft_export_report_dde.png
 ```
 
 ## The returned Exported Program
@@ -251,12 +251,3 @@ and produce a runnable artifact. This optimized version can then be used for
 deployment. In parallel, we can utilize the report generated by draft-export to
 identify and fix `torch.export` errors that were encountered so that the
 original model can be directly traceable with `torch.export`.
-
-```{toctree}
-:caption: Additional Links
-:maxdepth: 1
-
-torch.compiler_fake_tensor
-torch.compiler_dynamic_shapes
-torch.compiler_aot_inductor
-```
diff --git a/docs/source/export/ir_spec.md b/docs/source/export/ir_spec.md
diff --git a/docs/source/export/programming_model.md b/docs/source/export/programming_model.md
@@ -1,4 +1,4 @@
-(export-programming-model)=
+(export.programming_model)=
 
 # torch.export Programming Model
 
@@ -15,7 +15,9 @@ on different inputs as long as they satisfy the same conditions.
 
 The basic output of {func}`torch.export.export` is a single graph of PyTorch
 operations, with associated metadata. The exact format of this output is
-covered in the {ref}`export.ir_spec`.
+covered in the {ref}`export IR spec <export.ir_spec>`.
+
+(non-strict-export)=
 
 ### Strict vs. Non-Strict Tracing
 
@@ -120,6 +122,9 @@ Whether a value is static or dynamic depends on its type:
 
   - There are dynamic variants for some primitive types (`SymInt`,
     `SymFloat`, `SymBool`). Typically users do not have to deal with them.
+  - Users can specify integer inputs as dynamic by specifying
+    a [dynamic shape](https://pytorch.org/docs/main/export.html#expressing-dynamism)
+    for it.
 
 - For Python *standard containers* (`list`, `tuple`, `dict`, `namedtuple`):
 
@@ -150,7 +155,7 @@ By default, the types of inputs you can use for your program are:
 - Python primitives (`int`, `float`, `bool`, `str`, `None`)
 - Python standard containers (`list`, `tuple`, `dict`, `namedtuple`)
 
-### Custom Input Types
+### Custom Input Types (PyTree)
 
 In addition, you can also define your own (custom) class and use it as an
 input type, but you will need to register such a class as a PyTree.
@@ -164,7 +169,8 @@ class Input:
     f: torch.Tensor
     p: torch.Tensor
 
-torch.export.register_dataclass(Input)
+import torch.utils._pytree as pytree
+pytree.register_dataclass(Input)
 
 class M(torch.nn.Module):
     def forward(self, x: Input):
diff --git a/docs/source/export/pt2_archive.md b/docs/source/export/pt2_archive.md
@@ -0,0 +1,122 @@
+(export.pt2_archive)=
+
+# PT2 Archive Spec
+
+The following specification defines the archive format which can be produced
+through the following methods:
+
+* {ref}`torch.export <torch.export>` through calling {func}`torch.export.save`
+* {ref}`AOTInductor <torch.compiler_aot_inductor>` through calling {func}`torch._inductor.aoti_compile_and_package`
+
+The archive is a zipfile, and can be manipulated using standard zipfile APIs.
+
+The following is a sample archive. We will walk through the archive folder by folder.
+
+```
+.
+├── archive_format
+├── byteorder
+├── .data
+│   ├── serialization_id
+│   └── version
+├── data
+│   ├── aotinductor
+│   │   └── model1
+│   │       ├── aotinductor_pickle_data.json
+│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.cpp
+│   │       ├── cf5ez6ifexr7i2hezzz4s7xfusj4wtisvu2gddeamh37bw6bghjw.so
+│   │       ├── cg7domx3woam3nnliwud7yvtcencqctxkvvcafuriladwxw4nfiv.cubin
+│   │       └── cubaaxppb6xmuqdm4bej55h2pftbce3bjyyvljxbtdfuolmv45ex.cubin
+│   ├── weights
+│   │  ├── model1_model_param_config.json
+│   │  ├── weight_0
+│   │  ├── weight_1
+│   │  ├── weight_2
+│   └── constants
+│   │  ├── model1_model_constants_config.json
+│   │  ├── tensor_0
+│   │  ├── tensor_1
+│   │  ├── custom_obj_0
+│   │  ├── custom_obj_1
+│   └── sample_inputs
+│       ├── model1.pt
+│       └── model2.pt
+├── extra
+│   └── ....json
+└── models
+    ├── model1.json
+    └── model2.json
+```
+
+## Contents
+
+### Archive Headers
+
+* `archive_format` declares the format used by this archive. Currently, it can only be “pt2”.
+* `byteorder`. One of “little” or “big”, used by zip file reader
+* `/.data/version` contains the archive version. (Notice that this is neither export serialization’s schema version, nor Aten Opset Version).
+* `/.data/serialization_id` is a hash generated for the current archive, used for verification.
+
+
+### AOTInductor Compiled Artifact
+
+Path: `/data/aotinductor/<model_name>-<backend>/`
+
+AOTInductor compilation artifacts are saved for each model-backend pair. For
+example, compilation artifacts for the `model1` model on A100 and H100 will be
+saved in `model1-a100` and `model1-h100` folders separately.
+
+The folder typically contains
+* `<uuid>.so`: Dynamic library compiled from <uuid>.cpp.
+* `<uuid>.cpp`: AOTInductor generated cpp wrapper file.
+* `*.cubin`: Triton kernels compiled from triton codegen kernels
+* (optional) `<uuid>.json`: External fallback nodes for custom ops to be executed by `ProxyExecutor`, serialized according to `ExternKernelNode` struct. If the model doesn’t use custom ops/ProxyExecutor, this file would be omitted.
+* `<uuid>_metadata.json`: Metadata which was passed in from the `aot_inductor.metadata` inductor config
+
+### Weights
+
+Path: `/data/weights/*`
+
+Model parameters and buffers are saved in the `/data/weights/` folder. Each
+tensor is saved as a separated file. The file only contains the raw data blob,
+tensor metadata are saved separately in the
+`<model_name>_model_param_config.json`.
+
+### Constants
+
+Path: `/data/constants/*`
+
+TensorConstants, non-persistent buffers and TorchBind objects are saved in the
+`/data/constants/` folder. Metadata is saved separately in the
+`<model_name>_model_constants_config.json`
+
+### Sample Inputs
+
+Path: `/data/sample_inputs/<model_name>.pt`
+
+The `sample_input` used by `torch.export` could be included in the archive for
+downstream use. Typically, it’s a flattened list of Tensors, combining both args
+and kwargs of the forward() function.
+
+The .pt file is produced by `torch.save(sample_input)`, and can be loaded by
+`torch.load()` in python and `torch::pickle_load()` in c++.
+
+When the model has multiple copies of sample input, it would be packaged as
+`<model_name>_<index>.pt`.
+
+### Models Definitions
+
+Path: `/models/<model_name>.json`
+
+Model definition is the serialized json of the ExportedProgram from
+`torch.export.save`, and other model-level metadata.
+
+## Multiple Models
+
+This archive spec supports multiple model definitions coexisting in the same
+file, with `<model_name>` serving as a unique identifier for the models, and
+will be used as reference in other folders of the archive.
+
+Lower level APIs like {func}`torch.export.pt2_archive._package.package_pt2` and
+{func}`torch.export.pt2_archive._package.load_pt2` allow you to have
+finer-grained control over the packaging and loading process.
diff --git a/docs/source/torch.compiler_aot_inductor.md b/docs/source/torch.compiler_aot_inductor.md
@@ -1,3 +1,5 @@
+(torch.compiler_aot_inductor)=
+
 # AOTInductor: Ahead-Of-Time Compilation for Torch.Export-ed Models
 
 ```{warning}
@@ -25,7 +27,7 @@ relies on.
 
 We will then use {func}`torch._inductor.aoti_compile_and_package` to compile the
 exported program using TorchInductor, and save the compiled artifacts into one
-package.
+package. The package is in the format of a {ref}`PT2 Archive Spec <export.pt2_archive>`.
 
 ```{note}
 If you have a CUDA-enabled device on your machine and you installed PyTorch with CUDA support,
diff --git a/docs/source/torch.compiler_ir.md b/docs/source/torch.compiler_ir.md
@@ -1,3 +1,5 @@
+(torch.compiler_ir)=
+
 # IRs
 
 PyTorch 2.0 offers two set of IRs for backends to interface with: Core Aten IR and Prims IR.
diff --git a/torch/export/dynamic_shapes.py b/torch/export/dynamic_shapes.py
@@ -85,15 +85,19 @@ def __call__(self, min=None, max=None) -> "_DimHint":
 
 class Dim:
     """
-    The `Dim` class allows users to specify dynamism in their exported programs. By marking a dimension with a `Dim`,
-    the compiler associates the dimension with a symbolic integer containing a dynamic range.
+    The ``Dim`` class allows users to specify dynamism in their exported
+    programs. By marking a dimension with a ``Dim``, the compiler associates the
+    dimension with a symbolic integer containing a dynamic range.
 
-    The API can be used in 2 ways: Dim hints (i.e. automatic dynamic shapes: `Dim.AUTO`, `Dim.DYNAMIC`, `Dim.STATIC`),
-    or named Dims (i.e. `Dim("name", min=1, max=2)`).
+    The API can be used in 2 ways: Dim hints (i.e. automatic dynamic shapes:
+    ``Dim.AUTO``, ``Dim.DYNAMIC``, ``Dim.STATIC``), or named Dims (i.e.
+    ``Dim("name", min=1, max=2)``).
 
-    Dim hints provide the lowest barrier to exportability, with the user only needing to specify if a dimension
-    if dynamic, static, or left for the compiler to decide (`Dim.AUTO`). The export process will automatically
-    infer the remaining constraints on min/max ranges and relationships between dimensions.
+    Dim hints provide the lowest barrier to exportability, with the user only
+    needing to specify if a dimension if dynamic, static, or left for the
+    compiler to decide (``Dim.AUTO``). The export process will automatically
+    infer the remaining constraints on min/max ranges and relationships between
+    dimensions.
 
     Example::
 
@@ -112,19 +116,19 @@ def forward(self, x, y):
         }
         ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)
 
-    Here, export would raise an exception if we replaced all uses of `Dim.AUTO` with `Dim.DYNAMIC`,
-    as x.shape[0] is constrained to be static by the model.
+    Here, export would raise an exception if we replaced all uses of ``Dim.AUTO`` with ``Dim.DYNAMIC``,
+    as ``x.shape[0]`` is constrained to be static by the model.
 
     More complex relations between dimensions may also be codegened as runtime assertion nodes by the compiler,
-    e.g. (x.shape[0] + y.shape[1]) % 4 == 0, to be raised if runtime inputs do not satisfy such constraints.
+    e.g. ``(x.shape[0] + y.shape[1]) % 4 == 0``, to be raised if runtime inputs do not satisfy such constraints.
 
-    You may also specify min-max bounds for Dim hints, e.g. `Dim.AUTO(min=16, max=32)`, `Dim.DYNAMIC(max=64)`,
+    You may also specify min-max bounds for Dim hints, e.g. ``Dim.AUTO(min=16, max=32)``, ``Dim.DYNAMIC(max=64)``,
     with the compiler inferring the remaining constraints within the ranges. An exception will be raised if
     the valid range is entirely outside the user-specified range.
 
     Named Dims provide a stricter way of specifying dynamism, where exceptions are raised if the compiler
     infers constraints that do not match the user specification. For example, exporting the previous
-    model, the user would need the following `dynamic_shapes` argument::
+    model, the user would need the following ``dynamic_shapes`` argument::
 
         s0 = Dim("s0")
         s1 = Dim("s1", min=16)
@@ -134,8 +138,9 @@ def forward(self, x, y):
         }
         ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)
 
-    Named Dims also allow specification of relationships between dimensions, up to univariate linear relations.
-    For example, the following indicates one dimension is a multiple of another plus 4::
+    Named Dims also allow specification of relationships between dimensions, up
+    to univariate linear relations.  For example, the following indicates one
+    dimension is a multiple of another plus 4::
 
         s0 = Dim("s0")
         s1 = 3 * s0 + 4
diff --git a/torch/export/pt2_archive/_package.py b/torch/export/pt2_archive/_package.py
@@ -12,8 +12,8 @@
 import torch
 import torch.utils._pytree as pytree
 from torch._export.serde.serialize import deserialize, serialize, SerializedArtifact
+from torch.export import ExportedProgram
 from torch.export._tree_utils import reorder_kwargs
-from torch.export.exported_program import ExportedProgram
 from torch.export.pt2_archive._package_weights import (
     get_complete,
     group_weights,
@@ -350,22 +350,21 @@ def package_pt2(
     opset_version: Optional[dict[str, int]] = None,
     pickle_protocol: int = DEFAULT_PICKLE_PROTOCOL,
 ) -> FileLike:
-    """
-    Saves the artifacts to a PT2Archive format
-    (https://docs.google.com/document/d/1RQ4cmywilnFUT1VE-4oTGxwXdc8vowCSZsrRgo3wFA8/edit?tab=t.0#heading=h.v2y2jgnwc56a).
-    The artifact can then be loaded using ``load_pt2``.
+    r"""
+    Saves the artifacts to a PT2Archive format. The artifact can then be loaded
+    using ``load_pt2``.
 
     Args:
-        f (str | os.PathLike[str] | IO[bytes]) A file-like object (has to
+        f (str | os.PathLike[str] | IO[bytes]): A file-like object (has to
          implement write and flush) or a string containing a file name.
 
         exported_programs (Union[ExportedProgram, dict[str, ExportedProgram]]):
          The exported program to save, or a dictionary mapping model name to an
          exported program to save. The exported program will be saved under
-         models/*.json. If only one ExportedProgram is specified, this will
+         models/\*.json. If only one ExportedProgram is specified, this will
          automatically be named "model".
 
-        aoti_files (Union[list[str], dict[str, list[str]]): A list of files
+        aoti_files (Union[list[str], dict[str, list[str]]]): A list of files
          generated by AOTInductor via
          ``torch._inductor.aot_compile(..., {"aot_inductor.package": True})``,
          or a dictionary mapping model name to its AOTInductor generated files.
diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py
@@ -15,10 +15,10 @@
 import torch.fx._pytree as fx_pytree
 import torch.utils._pytree as pytree
 from torch._library.fake_class_registry import FakeScriptObject
+from torch.export import ExportedProgram
 from torch.export._tree_utils import reorder_kwargs
 from torch.export.exported_program import (
     ConstantArgument,
-    ExportedProgram,
     ExportGraphSignature,
     InputKind,
     ModuleCallSignature,

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+(torch.compiler_ir)=`
	`2`	`+`
`1`	`3`	`# IRs`
`2`	`4`
`3`	`5`	`PyTorch 2.0 offers two set of IRs for backends to interface with: Core Aten IR and Prims IR.`