pytorch
2f04a852 - [torch export] add source_fn_stack metadata in non-strict mode (#121909)

Commit
2 years ago
[torch export] add source_fn_stack metadata in non-strict mode (#121909) Summary: - add ProxyTorchFunctionMode to capture the high level torch functions - save the func info into tracer.traced_func, which can be retrieved later in ProxyTorchDispatchMode when we construct and populate meta for proxy nodes - need to add some filtering to bypass `if func in torch_function_passthrough or func in all_prims() or func in ProxyTorchFunctionMode.metadata_fns:` otherwise funcs invoked in tracing will also be captured - moved `PythonKeyTracer` more up as it's used in a couple of places before it's currently declared after I added some type annotations Test Plan: ``` buck2 run fbcode//mode/dev-nosan fbcode//caffe2/test:test_export -- -r test_non_strict_source_fn_stack ``` Differential Revision: D54873230
Parents
Loading