IvanKobzarev 8ae4f82243 [aotd] Support HOP effects in backward (#132638) ... Support of effectful operations in backward: 1/ AOTD collects metadata from forward fn only, so we can have usage of effectful ops in backward, that were not used in forward => Allowing tokens discovery during joint function . FunctionalTensorMode holds _tokens, in Joint function after tracing forward we memoize _tokens as `_tokens_forward_output`. 2/ Tokens are added as primals inputs (forward) in EffectTokensWrapper. Tokens that will be used in backward are in partitioner saved values. We do not have control on which positions they are saved in forward outputs. 2/ If new tokens discovered in backward after tracing joint_fn, the result graph will be manually added in the end of primals. _aot_autograd/utils.py 3/ All effectful ops during backward are marked with 'must_be_in_backward' partitioner_tag, to prevent partiitoner to place them in forward. For that functional_tensor_mode got new optional state `self._effects_partitioner_tag` for effectful ops, to set after tracing forward. There are additional changes in partitioner to improve functionality of 'must_be_in_backward' 4/ Unlift tokens now should run for both forward and backward. - As saved for backward tokens are placed on non static places - we identify input and output tokens to erase, by input and output of `with_effects` operation - In forward we can have input tokens, discovered in backward, that are not used in with_effects ops in forward, but saved for backward. We identify them by position in forward inputs. 5/ Adding aot debug logging for graphs before unlifting and before adding additional primal for backward tokens. Tests: ``` python test/higher_order_ops/test_with_effects.py ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/132638 Approved by: https://github.com/bdhirsh		2024-08-23 15:30:58 +00:00
..
_aot_autograd	[aotd] Support HOP effects in backward (#132638)	2024-08-23 15:30:58 +00:00
__init__.py
aot_autograd.py	Run cudagraphs on AOTAutograd cache hit (#132294)	2024-08-17 21:24:54 +00:00
apis.py	Revert "Deprecate torch._utils.is_compiling() and torch._dynamo.external_utils.is_compiling() (#127690)"	2024-08-07 00:05:20 +00:00
autograd_function.py	Add None return type to init (#132335)	2024-08-01 15:26:45 +00:00
batch_norm_replacement.py	Revert "[BE] typing for decorators - _library/custom_ops (#131578)"	2024-07-28 03:29:32 +00:00
benchmark_utils.py	[BE]: Apply ruff FURB 118. (#124743)	2024-04-26 14:34:52 +00:00
compile_utils.py	make functorch CSE respect mutations as barriers (like fsdp.set_) (#132243)	2024-08-05 21:28:55 +00:00
compilers.py	[BE] UFMT directory torch/_functorch (#123723)	2024-04-12 08:04:51 +00:00
config.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
deprecated.py	Flip default value for mypy disallow_untyped_defs [2/11] (#127839)	2024-06-08 18:23:08 +00:00
eager_transforms.py	Fix linearize(grad(...)) call (#133364)	2024-08-15 17:55:36 +00:00
functional_call.py	Let dynamo inline functional_call (#128646)	2024-07-30 14:22:23 +00:00
fx_minifier.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
make_functional.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
partitioners.py	[aotd] Support HOP effects in backward (#132638)	2024-08-23 15:30:58 +00:00
pyfunctorch.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
python_key.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
pytree_hacks.py	[BE] wrap deprecated function/class with typing_extensions.deprecated (#127689)	2024-06-02 12:30:43 +00:00
top_operators_github_usage.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
utils.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00
vmap.py	[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)	2024-07-31 10:42:50 +00:00