to_onnx return ONNXProgram

pyproject.toml

@@ -180,6 +180,7 @@ markers = [
 ]
 filterwarnings = [
     "error::FutureWarning",
+    "ignore::FutureWarning:onnxscript", # Temporary ignore until onnxscript is updated
 ]
 xfail_strict = true
 junit_duration_report = "call"

requirements/pytorch/test.txt

@@ -11,6 +11,7 @@ scikit-learn >0.22.1, <1.7.0
 numpy >=1.17.2, <1.27.0
 onnx >=1.12.0, <1.18.0
 onnxruntime >=1.12.0, <1.21.0
+onnxscript >= 0.2.2, <0.2.6
 psutil <7.0.1 # for `DeviceStatsMonitor`
 pandas >1.0, <2.3.0  # needed in benchmarks
 fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in App

src/lightning/pytorch/core/module.py

@@ -74,8 +74,10 @@
 
 if TYPE_CHECKING:
     from torch.distributed.device_mesh import DeviceMesh
+    from torch.onnx import ONNXProgram
 
 _ONNX_AVAILABLE = RequirementCache("onnx")
+_ONNXSCRIPT_AVAILABLE = RequirementCache("onnxscript")
 
 warning_cache = WarningCache()
 log = logging.getLogger(__name__)
@@ -1360,12 +1362,18 @@ def _verify_is_manual_optimization(self, fn_name: str) -> None:
             )
 
     @torch.no_grad()
-    def to_onnx(self, file_path: Union[str, Path, BytesIO], input_sample: Optional[Any] = None, **kwargs: Any) -> None:
+    def to_onnx(
+        self,
+        file_path: Union[str, Path, BytesIO, None] = None,
+        input_sample: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> Union["ONNXProgram", None]:
         """Saves the model in ONNX format.
 
         Args:
-            file_path: The path of the file the onnx model should be saved to.
+            file_path: The path of the file the onnx model should be saved to. Default: None (no file saved).
             input_sample: An input for tracing. Default: None (Use self.example_input_array)
+
             **kwargs: Will be passed to torch.onnx.export function.
 
         Example::
@@ -1386,6 +1394,11 @@ def forward(self, x):
         if not _ONNX_AVAILABLE:
             raise ModuleNotFoundError(f"`{type(self).__name__}.to_onnx()` requires `onnx` to be installed.")
 
+        if kwargs.get("dynamo", False) and not _ONNXSCRIPT_AVAILABLE:
+            raise ModuleNotFoundError(
+                f"`{type(self).__name__}.to_onnx(dynamo=True)` requires `onnxscript` to be installed."
+            )
+
         mode = self.training
 
         if input_sample is None:
@@ -1402,8 +1415,9 @@ def forward(self, x):
         file_path = str(file_path) if isinstance(file_path, Path) else file_path
         # PyTorch (2.5) declares file_path to be str | PathLike[Any] | None, but
         #               BytesIO does work, too.
-        torch.onnx.export(self, input_sample, file_path, **kwargs)  # type: ignore
+        ret = torch.onnx.export(self, input_sample, file_path, **kwargs)  # type: ignore
         self.train(mode)
+        return ret
 
     @torch.no_grad()
     def to_torchscript(

src/lightning/pytorch/utilities/testing/_runif.py

@@ -18,7 +18,7 @@
 from lightning.fabric.utilities.testing import _runif_reasons as fabric_run_if
 from lightning.pytorch.accelerators.cpu import _PSUTIL_AVAILABLE
 from lightning.pytorch.callbacks.progress.rich_progress import _RICH_AVAILABLE
-from lightning.pytorch.core.module import _ONNX_AVAILABLE
+from lightning.pytorch.core.module import _ONNX_AVAILABLE, _ONNXSCRIPT_AVAILABLE
 from lightning.pytorch.utilities.imports import _OMEGACONF_AVAILABLE
 
 _SKLEARN_AVAILABLE = RequirementCache("scikit-learn")
@@ -42,6 +42,7 @@ def _runif_reasons(
     psutil: bool = False,
     sklearn: bool = False,
     onnx: bool = False,
+    onnxscript: bool = False,
 ) -> tuple[list[str], dict[str, bool]]:
     """Construct reasons for pytest skipif.
 
@@ -64,6 +65,7 @@ def _runif_reasons(
         psutil: Require that psutil is installed.
         sklearn: Require that scikit-learn is installed.
         onnx: Require that onnx is installed.
+        onnxscript: Require that onnxscript is installed.
 
     """
 
@@ -96,4 +98,7 @@ def _runif_reasons(
     if onnx and not _ONNX_AVAILABLE:
         reasons.append("onnx")
 
+    if onnxscript and not _ONNXSCRIPT_AVAILABLE:
+        reasons.append("onnxscript")
+
     return reasons, kwargs

tests/tests_pytorch/models/test_onnx.py

@@ -167,3 +167,13 @@ def to_numpy(tensor):
 
     # compare ONNX Runtime and PyTorch results
     assert np.allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
+
+
+@RunIf(onnx=True, min_torch="2.7.0", dynamo=True, onnxscript=True)
+def test_model_return_type():
+    model = BoringModel()
+    model.example_input_array = torch.randn((1, 32))
+    model.eval()
+
+    ret = model.to_onnx(dynamo=True)
+    assert isinstance(ret, torch.onnx.ONNXProgram)