mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
8d921eb97f
The `export` API takes a `nn.Module` and traces its `forward` method. However sometimes it is useful to export different methods of a `nn.Module`, either as a one-off for debugging or as a set of methods that are called in some sequence outside `export` (e.g., `encode` / `decode`). When multiple methods of the same module instance are exported, they should share the same of the common module instance. This PR adds a couple of utils in `torch._export.utils` for this workflow. The `wrap_method` util wraps a method as a `nn.Module` that can then be exported. See included test. We recommend using the same module instance to export multiple methods on that instance, in which case they are guaranteed to share state. On serde, this state sharing is lost, so we provide another util, `sync_state`, to re-sync the state. These utils are meant to be eventually replaced by API-level changes, but for now this can unblock users who need this workflow. In particular, in the future we can accept one or multiple method entrypoints, with their own args / kwargs / dynamic shape specifications, which can create a variant of `ExportedProgram` with multiple graphs that share state; then we can automatically ensure that the state sharing is preserved through serde. Differential Revision: [D69960801](https://our.internmc.facebook.com/intern/diff/D69960801/) Pull Request resolved: https://github.com/pytorch/pytorch/pull/147573 Approved by: https://github.com/tugsbayasgalan |
||
|---|---|---|
| .. | ||
| db | ||
| pass_infra | ||
| passes | ||
| serde | ||
| __init__.py | ||
| converter.py | ||
| error.py | ||
| non_strict_utils.py | ||
| pass_base.py | ||
| tools.py | ||
| utils.py | ||
| verifier.py | ||
| wrappers.py | ||