frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Docs on export joint with descriptors (#159006)

Browse Source 
Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/159006
Approved by: https://github.com/SherlockNoMad
This commit is contained in:
Edward Z. Yang
2025-08-11 21:38:02 -04:00
committed by PyTorch MergeBot
parent 20629b1619
commit b2b4add0e7
2 changed files with 112 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
docs/source/export.md
View File

@ -645,6 +645,7 @@ export/programming_model
export/ir_spec
export/pt2_archive
export/draft_export
export/joint_with_descriptors
cond
generated/exportdb/index
torch.compiler_aot_inductor

111
docs/source/export/joint_with_descriptors.md Normal file
View File

@ -0,0 +1,111 @@
# Joint with descriptors
Joint with descriptors is an experimental API for exporting a traced joint
graph that supports all of torch.compile's features in full generality and,
after processing, can be converted back into a differentiable callable that
can be executed as normal. For example, it is used to implement autoparallel,
a system that takes a model and reshards inputs and parameters to make it
a distributed SPMD program.
```{eval-rst}
.. currentmodule:: torch._functorch.aot_autograd
.. autofunction:: aot_export_joint_with_descriptors
.. autofunction:: aot_compile_joint_with_descriptors
```
## Descriptors
```{eval-rst}
.. currentmodule:: torch._functorch._aot_autograd.descriptors
.. autoclass:: AOTInput
:members:
.. autoclass:: AOTOutput
:members:
.. autoclass:: BackwardTokenAOTInput
:members:
.. autoclass:: BackwardTokenAOTOutput
:members:
.. autoclass:: BufferAOTInput
:members:
.. autoclass:: DummyAOTInput
:members:
.. autoclass:: DummyAOTOutput
:members:
.. autoclass:: GradAOTOutput
:members:
.. autoclass:: InputMutationAOTOutput
:members:
.. autoclass:: IntermediateBaseAOTOutput
:members:
.. autoclass:: ParamAOTInput
:members:
.. autoclass:: PhiloxBackwardBaseOffsetAOTInput
:members:
.. autoclass:: PhiloxBackwardSeedAOTInput
:members:
.. autoclass:: PhiloxForwardBaseOffsetAOTInput
:members:
.. autoclass:: PhiloxForwardSeedAOTInput
:members:
.. autoclass:: PhiloxUpdatedBackwardOffsetAOTOutput
:members:
.. autoclass:: PhiloxUpdatedForwardOffsetAOTOutput
:members:
.. autoclass:: PlainAOTInput
:members:
.. autoclass:: PlainAOTOutput
:members:
.. autoclass:: SavedForBackwardsAOTOutput
:members:
.. autoclass:: SubclassGetAttrAOTInput
:members:
.. autoclass:: SubclassGetAttrAOTOutput
:members:
.. autoclass:: SubclassSizeAOTInput
:members:
.. autoclass:: SubclassSizeAOTOutput
:members:
.. autoclass:: SubclassStrideAOTInput
:members:
.. autoclass:: SubclassStrideAOTOutput
:members:
.. autoclass:: SyntheticBaseAOTInput
:members:
.. autoclass:: ViewBaseAOTInput
:members:
```
## FX utilities
```{eval-rst}
.. automodule:: torch._functorch._aot_autograd.fx_utils
:members:
```