pytorch/torch/_functorch
Avik Chaudhuri 3ed90f5a09 outline various stages from aot stage2 compile (#164808)
Splits the training and inference paths for aot stage2 compile.
1. Split `aot_stage2_autograd` into `_aot_stage2a_partition`, `_aot_stage2b_fw_compile` and `_aot_stage2b_bw_compile`, and rest.
2. Split `aot_stage2_inference` into `_aot_stage2b_inference_compile` and rest.
I'm leaving these as functions with underscore names since the I/O interfaces and the exact boundaries of these splits are somewhat in the air.

Differential Revision: D84028203

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164808
Approved by: https://github.com/SherlockNoMad
2025-10-10 17:04:36 +00:00
..
_activation_checkpointing [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
_aot_autograd outline various stages from aot stage2 compile (#164808) 2025-10-10 17:04:36 +00:00
__init__.py
aot_autograd.py move partition and compiler fns from stage 1 to stage 2 (#164765) 2025-10-07 18:02:03 +00:00
apis.py [BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551) 2025-06-25 06:16:06 +00:00
autograd_function.py Add pyrefly suppressions (3/n) (#164588) 2025-10-03 22:03:03 +00:00
batch_norm_replacement.py
benchmark_utils.py [BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551) 2025-06-25 06:16:06 +00:00
compile_utils.py [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
compilers.py [1/N] Use "is" in python type comparison (#165037) 2025-10-10 12:36:50 +00:00
config.py Revert "[2/N] More ruff SIM fixes (#165031)" 2025-10-10 13:42:14 +00:00
deprecated.py [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
eager_transforms.py [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
functional_call.py [2/N] Use "is" in python type comparison (#165142) 2025-10-10 15:36:44 +00:00
fx_minifier.py [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
make_functional.py [2/N] Use "is" in python type comparison (#165142) 2025-10-10 15:36:44 +00:00
partitioners.py Revert "[2/N] More ruff SIM fixes (#165031)" 2025-10-10 13:42:14 +00:00
predispatch.py Add support for tracing vmap in pre-dispatch export (#154650) 2025-08-20 19:31:07 +00:00
pyfunctorch.py Add pyrefly suppressions (3/n) (#164588) 2025-10-03 22:03:03 +00:00
python_key.py
pytree_hacks.py
top_operators_github_usage.py [BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551) 2025-06-25 06:16:06 +00:00
utils.py PEP585 update - torch/_functorch (#145139) 2025-01-19 07:06:10 +00:00
vmap.py [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00