[ONNX] Enable experimental exporter logic to dynamo_export and support refine dynamic_shapes (#134976)

(1) Enable experimental exporter logic to dynamo_export
(2) Refine dynamic shapes and retry export in export strategies
(3) Delete `torch_export_graph_extractor` and use the new export logic
(4) Disable ExportedProgram test in `test_fx_onnx_with_onnxruntime.py`, as ONNXProgram is different now.

Fixes https://github.com/pytorch/pytorch/issues/126479
Fixes #135183
Pull Request resolved: https://github.com/pytorch/pytorch/pull/134976
Approved by: https://github.com/justinchuby
This commit is contained in:
titaiwangms
2024-09-06 01:29:54 +00:00
committed by PyTorch MergeBot
parent 1e57ef08fa
commit 8f6e73f068
12 changed files with 246 additions and 411 deletions

View File

@ -346,6 +346,28 @@ def skipDtypeChecking(func):
return wrapper
def skip_if_fake_model_and_inititalizer(reason: Optional[str] = None):
"""skip test with models using ExportedProgram as input.
Args:
reason: The reason for skip the ONNX export test.
Returns:
A decorator for skip tests.
"""
def skip_dec(func):
@functools.wraps(func)
def wrapper(self, *args, **kwargs):
if kwargs["use_fake_mode"] and kwargs["include_initializer"]:
return unittest.SkipTest(reason)
return func(self, *args, **kwargs)
return wrapper
return skip_dec
def xfail_if_model_type_is_exportedprogram(
error_message: str, reason: Optional[str] = None
):