Rename _torch_onnx_export

justinchuby · Sep 6, 2024 · a8d7d3b · a8d7d3b
1 parent 0f5afaf
commit a8d7d3b
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/src/torch_onnx/__init__.py b/src/torch_onnx/__init__.py
@@ -17,6 +17,6 @@
 from ._analysis import analyze
 from ._core import export, exported_program_to_ir
 from ._onnx_program import ONNXProgram
-from ._patch import _export_compat as export_compat
+from ._patch import _torch_onnx_export as export_compat
 from ._patch import patch_torch, unpatch_torch
 from ._registration import ONNXRegistry
diff --git a/src/torch_onnx/_patch.py b/src/torch_onnx/_patch.py
@@ -117,7 +117,7 @@ def _get_torch_export_args(
     return args, kwargs
 
 
-def _export_compat(
+def _torch_onnx_export(
     model: torch.nn.Module | torch.export.ExportedProgram,
     args: tuple[Any, ...],
     f: str | os.PathLike | None = None,
@@ -261,7 +261,7 @@ def _to_dynamic_shape(x):
     else:
         dynamic_shapes = None
 
-    return _export_compat(
+    return _torch_onnx_export(
         model,
         model_args,
         kwargs=model_kwargs,
@@ -311,7 +311,7 @@ def patch_torch(
     _ARTIFACTS_DIR = artifacts_dir
     global _FALLBACK_TO_LEGACY_EXPORT  # noqa: PLW0603
     _FALLBACK_TO_LEGACY_EXPORT = fallback
-    torch.onnx.export = _export_compat  # type: ignore[assignment]
+    torch.onnx.export = _torch_onnx_export  # type: ignore[assignment]
     torch.onnx.dynamo_export = _torch_onnx_dynamo_export  # type: ignore[assignment]
 
 

diff --git a/tests/api_test.py b/tests/api_test.py
@@ -191,7 +191,7 @@ def forward(self, x, y):
             "x": (torch.export.Dim("dx"),),
             "y": (torch.export.Dim("dy"),),
         }
-        onnx_program = torch_onnx._patch._export_compat(
+        onnx_program = torch_onnx.export_compat(
             TestRefineDynamicShapeModel(), inps, dynamic_shapes=dynamic_shapes
         )
         assert onnx_program is not None