[BugFix] ONNX Export Syntax in Tests and Add Optional Dependencies (#1478)

Co-authored-by: Antoine de Maleprade <antoine.demaleprade@helsing.ai>

Author: az0uz

pyproject.toml

@@ -58,6 +58,7 @@ tests = [
 h5 = ["h5py>=3.8"]
 dev = ["pybind11", "ninja"]
 typecheck = ["mypy>=1.0.0"]
+onnx = ["onnx", "onnxscript", "onnxruntime"]
 
 [tool.setuptools]
 include-package-data = false

test/test_compile.py

@@ -950,9 +950,7 @@ def test_onnx_export_module(self, tmpdir):
         x = torch.randn(3)
         y = torch.randn(3)
         torch_input = {"x": x, "y": y}
-        onnx_program = torch.onnx.dynamo_export(tdm, **torch_input)
-
-        onnx_input = onnx_program.adapt_torch_inputs_to_onnx(**torch_input)
+        onnx_program = torch.onnx.export(tdm, kwargs=torch_input, dynamo=True)
 
         path = Path(tmpdir) / "file.onnx"
         onnx_program.save(str(path))
@@ -969,9 +967,7 @@ def to_numpy(tensor):
                 else tensor.cpu().numpy()
             )
 
-        onnxruntime_input = {
-            k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)
-        }
+        onnxruntime_input = {k: to_numpy(v) for k, v in torch_input.items()}
 
         onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
         torch.testing.assert_close(
@@ -986,10 +982,8 @@ def test_onnx_export_seq(self, tmpdir):
         x = torch.randn(3)
         y = torch.randn(3)
         torch_input = {"x": x, "y": y}
-        torch.onnx.dynamo_export(tdm, x=x, y=y)
-        onnx_program = torch.onnx.dynamo_export(tdm, **torch_input)
-
-        onnx_input = onnx_program.adapt_torch_inputs_to_onnx(**torch_input)
+        torch.onnx.export(tdm, kwargs=torch_input, dynamo=True)
+        onnx_program = torch.onnx.export(tdm, kwargs=torch_input, dynamo=True)
 
         path = Path(tmpdir) / "file.onnx"
         onnx_program.save(str(path))
@@ -1006,9 +1000,7 @@ def to_numpy(tensor):
                 else tensor.cpu().numpy()
             )
 
-        onnxruntime_input = {
-            k.name: to_numpy(v) for k, v in zip(ort_session.get_inputs(), onnx_input)
-        }
+        onnxruntime_input = {k: to_numpy(v) for k, v in torch_input.items()}
 
         onnxruntime_outputs = ort_session.run(None, onnxruntime_input)
         torch.testing.assert_close(