fix

Author: xadupre

_unittests/ut_torch_onnx/test_sbs.py

@@ -1,4 +1,5 @@
 import unittest
+import onnx
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
     hide_stdout,
@@ -9,11 +10,7 @@
 from onnx_diagnostic.reference import ExtendedReferenceEvaluator, OnnxruntimeEvaluator
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
 from onnx_diagnostic.torch_onnx.sbs import run_aligned
-
-try:
-    from experimental_experiment.torch_interpreter import to_onnx
-except ImportError:
-    to_onnx = None
+from onnx_diagnostic.export.api import to_onnx
 
 
 class TestSideBySide(ExtTestCase):
@@ -41,7 +38,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (x,), dynamic_shapes=({0: self.torch.export.Dim("batch")},)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -71,10 +68,12 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (x,), dynamic_shapes=({0: self.torch.export.Dim("batch")},)
         )
-        epo = self.torch.onnx.export(
-            ep, (x,), dynamic_shapes=({0: self.torch.export.Dim("batch")},), dynamo=True
-        )
-        onx = epo.model_proto
+        onx = to_onnx(
+            ep,
+            (x,),
+            dynamic_shapes=({0: self.torch.export.Dim("batch")},),
+            exporter="onnx-dynamo",
+        ).model_proto
         results = list(
             run_aligned(
                 ep,
@@ -105,9 +104,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        epo = self.torch.onnx.export(
-            Model(), (), kwargs=inputs, dynamic_shapes=ds, dynamo=True
-        )
+        epo = to_onnx(Model(), (), kwargs=inputs, dynamic_shapes=ds, exporter="onnx-dynamo")
         onx = epo.model_proto
         results = list(
             run_aligned(
@@ -139,7 +136,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -170,7 +167,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -204,7 +201,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -240,7 +237,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -275,7 +272,7 @@ def forward(self, x):
         ep = self.torch.export.export(
             Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
         )
-        onx = to_onnx(ep)
+        onx = to_onnx(ep, exporter="custom").model_proto
         results = list(
             run_aligned(
                 ep,
@@ -291,6 +288,45 @@ def forward(self, x):
         self.assertEqual(len(results), 7)
         self.assertEqual([r[-1].get("dev", 0) for r in results], [0, 0, 0, 0, 0, 0, 0])
 
+    @hide_stdout()
+    @ignore_warnings((DeprecationWarning, FutureWarning, UserWarning))
+    def test_sbs_model_with_weights(self):
+        torch = self.torch
+
+        class Model(self.torch.nn.Module):
+            def __init__(self):
+                super(Model, self).__init__()
+                self.fc1 = torch.nn.Linear(10, 32)  # input size 10 → hidden size 32
+                self.relu = torch.nn.ReLU()
+                self.fc2 = torch.nn.Linear(32, 1)  # hidden → output
+
+            def forward(self, x):
+                x = self.relu(self.fc1(x))
+                x = self.fc2(x)
+                return x
+
+        inputs = dict(x=self.torch.randn((5, 10)))
+        ds = dict(x={0: "batch"})
+        Model()(**inputs)
+        ep = self.torch.export.export(
+            Model(), (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds)
+        )
+        filename = self.get_dump_file("test_sbs_model_with_weights.onnx")
+        to_onnx(ep, exporter="custom", filename=filename)
+        onx = onnx.load(filename)
+        results = list(
+            run_aligned(
+                ep,
+                onx,
+                kwargs=inputs,
+                run_cls=OnnxruntimeEvaluator,
+                verbose=11,
+                use_tensor=True,
+            ),
+        )
+        self.assertEqual(len(results), 7)
+        self.assertEqual([r[-1].get("dev", 0) for r in results], [0, 0, 0, 0, 0, 0, 0])
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/_command_lines_parser.py

@@ -1230,7 +1230,7 @@ def post_process(obs):
     for obs in run_aligned(
         ep,
         onx,
-        run_cls=OnnxruntimeEvaluator,
+        run_cls=OnnxruntimeEvaluator,  # type: ignore[arg-type]
         atol=float(args.atol),
         rtol=float(args.rtol),
         verbose=int(args.verbose),

onnx_diagnostic/helpers/helper.py

@@ -1207,7 +1207,7 @@ def max_diff(
             if "dev" in d:
                 if dd is None:
                     dd = d["dev"]
-                else:
+                elif d["dev"] is not None:
                     dd += d["dev"]
 
         res = dict(abs=am, rel=rm, sum=sm, n=n, dnan=dn)
@@ -1264,7 +1264,7 @@ def max_diff(
             # out of boundary
             res = dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
             if dev:
-                res[dev] = dev
+                res["dev"] = dev
             return res
         if isinstance(expected, (int, float)):
             if isinstance(got, np.ndarray) and len(got.shape) == 0:
@@ -1280,7 +1280,7 @@ def max_diff(
                     dnan=0,
                 )
                 if dev:
-                    res[dev] = dev
+                    res["dev"] = dev
                 return res
             return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
         if expected.dtype in (np.complex64, np.complex128):
@@ -1362,7 +1362,7 @@ def max_diff(
             abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
         )
         if dev:
-            res[dev] = dev
+            res["dev"] = dev
         if hist:
             if isinstance(hist, bool):
                 hist = np.array([0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype)

onnx_diagnostic/reference/ort_evaluator.py

@@ -553,9 +553,13 @@ def _run(self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]) -> L
         feeds = dict(zip(node.input, inputs))
         if "" in feeds:
             cls = None
-            for k, v in feeds:
+            for k, v in feeds.items():
                 if k != "":
                     cls = v.__class__
+                    break
+            assert (
+                cls is not None
+            ), f"Unable to get input class (array or tensor), feeds={string_type(feeds)}"
             feeds[""] = cls([0])
 
         assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"

onnx_diagnostic/torch_onnx/sbs.py

@@ -537,9 +537,9 @@ def _loop_cmp(
         print(f"[run_aligned]   onnx {len(onnx_results)} constants")
         print(f"[run_aligned] common {len(mapping_onnx_to_torch)} constants")
         for k, v in torch_results.items():
-            print(f"[run_aligned-ep] +cst: {k}: {string_type(v, str_kws)}")
+            print(f"[run_aligned-ep] +cst: {k}: {string_type(v, **str_kws)}")
         for k, v in onnx_results.items():
-            print(f"[run_aligned-nx] +ini: {k}: {string_type(v, str_kws)}")
+            print(f"[run_aligned-nx] +ini: {k}: {string_type(v, **str_kws)}")
 
     onnx_args = list(args) if args else []
     if kwargs:
@@ -578,7 +578,7 @@ def _loop_cmp(
             if exc:
                 raise AssertionError(
                     f"unable to process node {node.op} -> {node.name!r}, "
-                    f"possible candiate are "
+                    f"possible candidate are "
                     f"{sorted(p for p in onnx_results if node.name in p)}, "
                     f"not in {sorted(onnx_results)}, "
                     f"args={string_type(args, **str_kws)}, "