Yidi Wu b9697eacd3 [torchbind] support tensor ops inside of __obj_flatten__ (#129605) ... As titled. Previously, __obj_flatten__ can run in a fake tensor mode, e.g. in process_input of aot_autograd, which is surrounded by a fake tensor mode. This causes the tensor ops inside __obj_flatten__ to run under fake tensor mode. However, tensors inside of script obejct are real tensors, this causes the fake tensor mode to error out saying that we need to first fakify fall the tensors (because allow_non_fake_inputs is set to True). In this PR, we disable all the dispatch modes when running to_fake_obj. Note that, the output of `__obj_flatten__` will be fakified and filled inside of the corresponding FakeScriptObject. So during traicng, we'll be using FakeScriptObject that has fake tensor contents. Test Plan: Add a new test: pytest test/export/test_torchbind.py -k test_compile_tensor_op_in_tensor_flatten Pull Request resolved: https://github.com/pytorch/pytorch/pull/129605 Approved by: https://github.com/angelayi		2024-06-27 03:07:31 +00:00
..
aoti_abi_check
aoti_inference
api	Remove more unused variables in tests (#127510)	2024-05-31 03:39:45 +00:00
c10d	TCPStore: improve connect and retry logic (#129261)	2024-06-25 19:24:22 +00:00
common
dist_autograd
jit	[torchbind] support tensor ops inside of __obj_flatten__ (#129605)	2024-06-27 03:07:31 +00:00
lazy	Remove unused functions (#127881)	2024-06-05 05:21:24 +00:00
lite_interpreter_runtime
monitor
profiler
rpc
tensorexpr	Check unused variables in tests (#127498)	2024-06-04 05:35:25 +00:00
__init__.py