Zhengxu Chen 53256edff9 [export] Support module inputs for non strict mode. (#143925) ... Summary: Add experimental support for torch.nn.Module as input types. Before this change, we don't support module inputs but recently we saw some interesting use cases like gpt-fast https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py#L68 where we directly pass in a module input for different variants of the same models. Since we don't really care about non-param or non-buffer states in non strict mode, we don't care about those either and pretend they are like plain constants during tracing. We treat any module input like a nested container of tensor, and each time we will automatically register a pytree handler for these module types to flatten its state dict into a group of tensors. We will just inline any module method call during tracing like we did for `self` module in export_for_training. This will make input modules' behavior very similar to the training module in typical case, except that we don't record the inputs as parameter or buffers but rather just plain user inputs. Test Plan: buck run mode/opt caffe2/test:test_export -- -r test_module_input Differential Revision: D67680827 Pull Request resolved: https://github.com/pytorch/pytorch/pull/143925 Approved by: https://github.com/tugsbayasgalan		2025-01-16 17:30:36 +00:00
..
__init__.py
opinfo_schema.py
random_dag.py	[BE] fix ruff rule E226: add missing whitespace around operator in f-strings (#144415)	2025-01-08 21:55:00 +00:00
test_converter.py
test_cpp_serdes.py
test_db.py
test_draft_export.py	[aoti] Remove example inputs from aoti_compile_and_package (#144520)	2025-01-10 21:56:23 +00:00
test_experimental.py	[export] handle buffer/input mutations for joint-graph (#144806)	2025-01-16 00:22:16 +00:00
test_export.py	[export] Support module inputs for non strict mode. (#143925)	2025-01-16 17:30:36 +00:00
test_export_legacy.py	EZ fix to make sure local pytest run succeeds in export (#144764)	2025-01-15 00:43:40 +00:00
test_export_nonstrict.py
test_export_training_ir_to_run_decomp.py
test_functionalized_assertions.py
test_hop.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_lift_unlift.py	[4/N] Apply py39 ruff and pyupgrade fixes (#143257)	2025-01-04 10:47:51 +00:00
test_pass_infra.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_passes.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_retraceability.py
test_schema.py
test_serdes.py
test_serialize.py	[export] Load side info about pos/kw argument kind for serialization. (#144686)	2025-01-15 19:08:38 +00:00
test_sparse.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_swap.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_tools.py
test_torchbind.py
test_tree_utils.py
test_unflatten.py	[Codemod][AddExplicitStrictExportArg] caffe2/test (#143688)	2024-12-27 07:58:44 +00:00
test_unflatten_training_ir.py
test_verifier.py
testing.py