Revert "[ONNX] New export logic leveraging ExportedProgram and ONNX IR (

pytorch#132530)"

This reverts commit 5fab35d.

Reverted pytorch#132530 on behalf of https://github.com/ZainRizvi due to Sorry but it seems like Dr. CI incorrectly flagged the [pull / linux-docs / build-docs-python-false](https://hud.pytorch.org/pr/pytorch/pytorch/132530#28918577682) failure as being flaky. The job started failing consistently on CI once your PR was merged. [GH job link](https://github.com/pytorch/pytorch/actions/runs/10454830880/job/28949386844) [HUD commit link](https://hud.pytorch.org/pytorch/pytorch/commit/5fab35d77c7d1db7dbb9d5c516254a510b4f4f64) ([comment](pytorch#132530 (comment)))

mypy.ini

@@ -165,6 +165,9 @@ ignore_missing_imports = True
 [mypy-tensorboard.*]
 ignore_missing_imports = True
 
+[mypy-onnx.*]
+ignore_missing_imports = True
+
 [mypy-matplotlib.*]
 ignore_missing_imports = True
 
@@ -298,14 +301,5 @@ ignore_missing_imports = True
 # Third party dependencies that are optional.
 #
 
-[mypy-onnx.*]
-ignore_missing_imports = True
-
-[mypy-onnxruntime.*]
-ignore_missing_imports = True
-
-[mypy-onnxscript.*]
-ignore_missing_imports = True
-
 [mypy-redis]
 ignore_missing_imports = True

test/onnx/dynamo/test_exporter_api.py

@@ -203,5 +203,222 @@ def test_serialize_succeeds_when_model_greater_than_2gb(self):
             serializer.serialize(onnx_program, io.BytesIO())
 
 
+class TestONNXExportWithDynamo(common_utils.TestCase):
+    def test_args_normalization_with_no_kwargs(self):
+        exported_program = torch.export.export(
+            SampleModelTwoInputs(),
+            (
+                torch.randn(1, 1, 2),
+                torch.randn(1, 1, 2),
+            ),
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program, torch.randn(1, 1, 2), torch.randn(1, 1, 2)
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), torch.randn(1, 1, 2)),
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_args_is_tensor_not_tuple(self):
+        exported_program = torch.export.export(SampleModel(), (torch.randn(1, 1, 2),))
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program, torch.randn(1, 1, 2)
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModel(), torch.randn(1, 1, 2), dynamo=True
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_args_normalization_with_kwargs(self):
+        exported_program = torch.export.export(
+            SampleModelTwoInputs(), (torch.randn(1, 1, 2),), {"b": torch.randn(1, 1, 2)}
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program, torch.randn(1, 1, 2), b=torch.randn(1, 1, 2)
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_args_normalization_with_empty_dict_at_the_tail(self):
+        exported_program = torch.export.export(
+            SampleModelTwoInputs(), (torch.randn(1, 1, 2),), {"b": torch.randn(1, 1, 2)}
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program, torch.randn(1, 1, 2), b=torch.randn(1, 1, 2)
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelTwoInputs(),
+            (torch.randn(1, 1, 2), {"b": torch.randn(1, 1, 2)}),
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_dynamic_axes_enable_dynamic_shapes_with_fully_specified_axes(self):
+        exported_program = torch.export.export(
+            SampleModelForDynamicShapes(),
+            (
+                torch.randn(2, 2, 3),
+                torch.randn(2, 2, 3),
+            ),
+            dynamic_shapes={
+                "x": {
+                    0: torch.export.Dim("customx_dim_0"),
+                    1: torch.export.Dim("customx_dim_1"),
+                    2: torch.export.Dim("customx_dim_2"),
+                },
+                "b": {
+                    0: torch.export.Dim("customb_dim_0"),
+                    1: torch.export.Dim("customb_dim_1"),
+                    2: torch.export.Dim("customb_dim_2"),
+                },
+            },
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program,
+            torch.randn(2, 2, 3),
+            b=torch.randn(2, 2, 3),
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "x": {0: "customx_dim_0", 1: "customx_dim_1", 2: "customx_dim_2"},
+                "b": {0: "customb_dim_0", 1: "customb_dim_1", 2: "customb_dim_2"},
+            },
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_dynamic_axes_enable_dynamic_shapes_with_default_axe_names(self):
+        exported_program = torch.export.export(
+            SampleModelForDynamicShapes(),
+            (
+                torch.randn(2, 2, 3),
+                torch.randn(2, 2, 3),
+            ),
+            dynamic_shapes={
+                "x": {
+                    0: torch.export.Dim("customx_dim_0"),
+                    1: torch.export.Dim("customx_dim_1"),
+                    2: torch.export.Dim("customx_dim_2"),
+                },
+                "b": {
+                    0: torch.export.Dim("customb_dim_0"),
+                    1: torch.export.Dim("customb_dim_1"),
+                    2: torch.export.Dim("customb_dim_2"),
+                },
+            },
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program,
+            torch.randn(2, 2, 3),
+            b=torch.randn(2, 2, 3),
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "x": [0, 1, 2],
+                "b": [0, 1, 2],
+            },
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_dynamic_axes_supports_partial_dynamic_shapes(self):
+        exported_program = torch.export.export(
+            SampleModelForDynamicShapes(),
+            (
+                torch.randn(2, 2, 3),
+                torch.randn(2, 2, 3),
+            ),
+            dynamic_shapes={
+                "x": None,
+                "b": {
+                    0: torch.export.Dim("customb_dim_0"),
+                    1: torch.export.Dim("customb_dim_1"),
+                    2: torch.export.Dim("customb_dim_2"),
+                },
+            },
+        )
+        onnx_program_from_new_exporter = torch.onnx.dynamo_export(
+            exported_program,
+            torch.randn(2, 2, 3),
+            b=torch.randn(2, 2, 3),
+        )
+        onnx_program_from_old_exporter = torch.onnx.export(
+            SampleModelForDynamicShapes(),
+            (torch.randn(2, 2, 3), {"b": torch.randn(2, 2, 3)}),
+            dynamic_axes={
+                "b": [0, 1, 2],
+            },
+            dynamo=True,
+        )
+        self.assertEqual(
+            onnx_program_from_new_exporter.model_proto,
+            onnx_program_from_old_exporter.model_proto,
+        )
+
+    def test_dynamic_shapes_hit_constraints_in_dynamo(self):
+        # SampleModelTwoInputs has constraints becuse of add of two inputs,
+        # so the two input shapes are related.
+        with self.assertRaisesRegex(
+            torch._dynamo.exc.UserError,
+            "Constraints violated",
+        ):
+            _ = torch.onnx.export(
+                SampleModelTwoInputs(),
+                (torch.randn(2, 2, 3), torch.randn(2, 2, 3)),
+                dynamic_axes={
+                    "x": {0: "x_dim_0", 1: "x_dim_1", 2: "x_dim_2"},
+                    "b": {0: "b_dim_0", 1: "b_dim_1", 2: "b_dim_2"},
+                },
+                dynamo=True,
+            )
+
+    def test_saved_f_exists_after_export(self):
+        with common_utils.TemporaryFileName(suffix=".onnx") as path:
+            _ = torch.onnx.export(
+                SampleModel(), torch.randn(1, 1, 2), path, dynamo=True
+            )
+            self.assertTrue(os.path.exists(path))
+
+    def test_raises_error_when_input_is_script_module(self):
+        class ScriptModule(torch.jit.ScriptModule):
+            def forward(self, x):
+                return x
+
+        with self.assertRaisesRegex(
+            TypeError,
+            "Dynamo export does not support ScriptModule or ScriptFunction.",
+        ):
+            _ = torch.onnx.export(ScriptModule(), torch.randn(1, 1, 2), dynamo=True)
+
+
 if __name__ == "__main__":
     common_utils.run_tests()

test/test_public_bindings.py

@@ -286,25 +286,6 @@ def onerror(modname):
         # do not get imported by public code.
         private_allowlist = {
             "torch._inductor.codegen.cuda.cuda_kernel",
-            # TODO(#133647): Remove the onnx._internal entries after
-            # onnx and onnxscript are installed in CI.
-            "torch.onnx._internal.exporter",
-            "torch.onnx._internal.exporter._analysis",
-            "torch.onnx._internal.exporter._building",
-            "torch.onnx._internal.exporter._capture_strategies",
-            "torch.onnx._internal.exporter._compat",
-            "torch.onnx._internal.exporter._core",
-            "torch.onnx._internal.exporter._decomp",
-            "torch.onnx._internal.exporter._dispatching",
-            "torch.onnx._internal.exporter._fx_passes",
-            "torch.onnx._internal.exporter._ir_passes",
-            "torch.onnx._internal.exporter._isolated",
-            "torch.onnx._internal.exporter._onnx_program",
-            "torch.onnx._internal.exporter._registration",
-            "torch.onnx._internal.exporter._reporting",
-            "torch.onnx._internal.exporter._schemas",
-            "torch.onnx._internal.exporter._tensors",
-            "torch.onnx._internal.exporter._verification",
             "torch.onnx._internal.fx._pass",
             "torch.onnx._internal.fx.analysis",
             "torch.onnx._internal.fx.analysis.unsupported_nodes",