[ONNX] Introduce Input/Ouptut formatter; Switch to 'DynamoExporter'

ghstack-source-id: c89f1cda4c217bb1cd593c13993eedf0c1e72b6a
Pull Request resolved: #98421

docs/source/onnx.rst

@@ -736,4 +736,4 @@ Preview: torch.onnx TorchDynamo Exporter
 
     torch.onnx.ExportOptions
     torch.onnx.ExportOutput
-    torch.onnx.ExportOutputSerializer
+    torch.onnx.ExportOutputSerializer

test/onnx/dynamo/test_exporter_api.py

@@ -8,6 +8,7 @@
 import torch
 from beartype import roar
 from torch.onnx import dynamo_export, ExportOptions, ExportOutput
+from torch.onnx._internal import exporter
 from torch.onnx._internal.exporter import (
     _DEFAULT_OPSET_VERSION,
     ExportOutputSerializer,
@@ -135,7 +136,9 @@ def serialize(
     def test_raise_on_invalid_save_argument_type(self):
         with self.assertRaises(roar.BeartypeException):
             ExportOutput(torch.nn.Linear(2, 3))  # type: ignore[arg-type]
-        export_output = ExportOutput(onnx.ModelProto())
+        export_output = ExportOutput(
+            onnx.ModelProto(), exporter.InputFormatter(), exporter.OutputFormatter()
+        )
         with self.assertRaises(roar.BeartypeException):
             export_output.save(None)  # type: ignore[arg-type]
         export_output.model_proto

test/onnx/pytorch_test_common.py

@@ -6,14 +6,14 @@
 import random
 import sys
 import unittest
-from typing import Optional
+from typing import Mapping, Optional, Type
 
 import numpy as np
 import packaging.version
 
 import torch
 from torch.autograd import function
-from torch.onnx._internal import diagnostics
+from torch.onnx._internal import diagnostics, exporter
 from torch.testing._internal import common_utils
 
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
@@ -206,6 +206,38 @@ def wrapper(self, *args, **kwargs):
     return skip_dec
 
 
+def skip_fx_exporters(
+    exporter_cls_and_reason: Mapping[Optional[Type[exporter.Exporter]], str]
+):
+    """Skip exporting test for selected FX exporters.
+
+    Args:
+        exporter_cls_and_reason: Mapping from FX exporter class to skip the test to the
+            reason for skipping.
+
+    Returns:
+        A decorator for skipping exporting test for FX exporters.
+    """
+
+    def skip_dec(func):
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            for exporter_cls, reason in exporter_cls_and_reason.items():
+                if exporter_cls == self.exporter_cls:
+                    exporter_name = (
+                        exporter_cls.__name__
+                        if exporter_cls is not None
+                        else "dynamo_export"
+                    )
+                    raise unittest.SkipTest(
+                        f"Skip verify test for '{exporter_name}'. {reason}"
+                    )
+
+        return wrapper
+
+    return skip_dec
+
+
 # skips tests for opset_versions listed in unsupported_opset_versions.
 # if the caffe2 test cannot be run for a specific version, add this wrapper
 # (for example, an op was modified but the change is not supported in caffe2)