Avik Chaudhuri 1d98a526dd preserve signatures with multiple calls + buffer mutations (#138669) ... As called out in https://github.com/pytorch/pytorch/pull/137999, preserving signatures of multiple calls when buffer mutations are present was NYI. The main problem was that intermediate values of buffers were not tracked, so couldn't be propagated statefully between multiple calls (i.e., they would need to be explicitly passed around, defeating the unlifting needed for preserving signatures). This PR fixes this situation, by introducing module attributes that carry the necessary intermediate values of buffer mutations. In general, a buffer mutation can have several intermediate values it depends on recursively, even other buffers. So rather than tying an intermediate value with a particular buffer, we tie it with the submodules that create and read it. We install an attribute on all modules that create or read a particular intermediate value, sharing the same initial storage (i.e., initialized with the same empty tensor). For the module that creates this intermediate value, we copy the value into the corresponding attribute; and for the modules that read it, we read the corresponding attribute instead. Another complication that needed to be addressed was that a `run_decompositions` following an `export_for_training` was not preserving module call graphs, which is needed for unflattening and, in particular, used when remapping inputs. Fortunately some existing metadata already tracks provenance of nodes, which we could use to update a module call graph after functionalization / decomposition. Differential Revision: D64806175 Pull Request resolved: https://github.com/pytorch/pytorch/pull/138669 Approved by: https://github.com/tugsbayasgalan		2024-10-25 00:13:25 +00:00
..
__init__.py
opinfo_schema.py
test_converter.py	[ts_converter] Fix prim::If buffer names (#136648)	2024-10-02 00:07:47 +00:00
test_db.py	Move test_db to training IR (#138251)	2024-10-18 07:42:13 +00:00
test_experimental.py	Move test_experimental.py to training IR (#138140)	2024-10-17 17:30:10 +00:00
test_export.py	preserve signatures with multiple calls + buffer mutations (#138669)	2024-10-25 00:13:25 +00:00
test_export_nonstrict.py
test_export_training_ir_to_run_decomp.py
test_functionalized_assertions.py
test_hop.py
test_lift_unlift.py
test_pass_infra.py
test_passes.py
test_retraceability.py	[export] Add retraceability_non_strict to tests (#138380)	2024-10-22 21:05:51 +00:00
test_schema.py
test_serdes.py	[export] Add serdes_non_strict to tests (#138662)	2024-10-24 21:35:32 +00:00
test_serialize.py	Move test_serialize to training IR (#138261)	2024-10-23 00:32:32 +00:00
test_sparse.py
test_swap.py	[export] Move swap to a different file (#137134)	2024-10-06 04:28:18 +00:00
test_tools.py
test_torchbind.py	[torchbind] Add generic __deepcopy__ method (#137613)	2024-10-24 22:14:55 +00:00
test_tree_utils.py
test_unflatten.py
test_unflatten_training_ir.py	fix unflatten training ir test suffix (#138840)	2024-10-24 23:42:54 +00:00
test_verifier.py	Move test_verifier to training IR (#138249)	2024-10-18 07:36:29 +00:00
testing.py	[export] Add serdes_non_strict to tests (#138662)	2024-10-24 21:35:32 +00:00