pytorch/test/package
Michael Suo 546db2e36e [fx] make fx.wrap idempotent (#104838)
Previously, if you called `torch.fx.wrap()` on the same thing twice, it would add two entries to `_wrapped_fns_to_patch`. Then, when tracing, the patcher would process them both. On the second entry, the patcher would double-wrap the fn (e.g. `wrap(wrap(orig_fn))`)

This makes it so that wrapping is observable after the trace. While normally, a Patcher instance will "revert" the wrapping after tracing, the double wrapped function goes from `wrap(wrap(orig_fn)) -> wrap(orig_fn)`.

This happens to work in normal fx stuff (after all, the wrapper function will behave exactly like the original function). But it upsets torch.package, which is not expecting to see a weird wrapper function in the graph.

This PR adds a dictionary to deduplicate `wrap()` calls, ensuring that the patcher only operates each once per frame-fn pair.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/104838
Approved by: https://github.com/Chillee
2023-07-09 20:57:46 +00:00
..
package_a
package_b
package_bc
package_c
package_d
package_e
test_trace_dep
__init__.py
common.py
generate_bc_packages.py
module_a.py
module_a_remapped_path.py
test_analyze.py
test_dependency_api.py
test_dependency_hooks.py
test_digraph.py
test_directory_reader.py
test_glob_group.py
test_importer.py
test_load_bc_packages.py
test_mangling.py
test_misc.py
test_model.py
test_package_fx.py
test_package_script.py
test_repackage.py
test_resources.py
test_save_load.py