onnx-diagnostic 0.8.3__py3-none-any.whl → 0.8.5__py3-none-any.whl

onnx_diagnostic/export/onnx_plug.py

@@ -3,8 +3,12 @@ from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import torch
-from ..helpers import max_diff
-from ..helpers.torch_helper import torch_dtype_to_onnx_dtype
+from ..helpers import max_diff, string_type
+from ..helpers.torch_helper import (
+    torch_dtype_to_onnx_dtype,
+    onnx_dtype_to_torch_dtype,
+    int_device_to_torch_device,
+)
 from ..reference import OnnxruntimeEvaluator
 
 TUPLE_TENSORS = Tuple[torch.Tensor, ...]
@@ -50,6 +54,10 @@ class EagerDirectReplacementWithOnnx:
         only tensors must be counted
     :param name: the name of the custom op, the function name if not specified
     :param kwargs: constants parameters with their default values
+    :param version_selector: selects the version based on the arguments,
+        see below for an example, this allows the user to define different
+        onnx version depending on the inputs
+    :param default_opset: opset to use by default
     :param verbose: verbose level
 
     Here is an example:
@@ -120,7 +128,61 @@ class EagerDirectReplacementWithOnnx:
 
         print(pretty_onnx(onx))
 
-        # And with :func:`torch.onnx.export`:
+    We do the same with :func:`torch.onnx.export`:
+
+    .. runpython::
+        :showcode:
+
+        import onnx.helper as oh
+        import torch
+        from onnx_diagnostic.helpers.onnx_helper import pretty_onnx
+        from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+        def demo_customsub(x, y):
+            return x - y
+
+
+        def demo_customsub_shape(x, y):
+            return torch.empty(torch.broadcast_shapes(x.shape, y.shape), dtype=x.dtype)
+
+
+        def make_function_proto():
+            return oh.make_function(
+                "onnx_plug",
+                "demo_customsub",
+                ["x", "y"],
+                ["z"],
+                [oh.make_node("Sub", ["x", "y"], ["z"])],
+                opset_imports=[oh.make_opsetid("", 22)],
+            )
+
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                y = x.sum(axis=1, keepdim=True)
+                d = torch.ops.onnx_plug.demo_customsub(x, y)
+                return torch.abs(d)
+
+
+        replacements = [
+            EagerDirectReplacementWithOnnx(
+                demo_customsub, demo_customsub_shape, make_function_proto(), 2, 1
+            )
+        ]
+
+        x = torch.randn((3, 4), dtype=torch.float32)
+        model = Model()
+        ds = ({0: "d1", 1: "d2"},)
+
+        # The exported program shows a custom op.
+        ep = torch.export.export(model, (x,), dynamic_shapes=use_dyn_not_str(ds))
+        print("ep")
+
+        # As the exporter knows how the replace this custom op.
+        # Let's export.
 
         onx = to_onnx(
             model,
@@ -133,27 +195,87 @@ class EagerDirectReplacementWithOnnx:
         ).model_proto
 
         print(pretty_onnx(onx))
+
+    This shows how to define multiple versions depending on the device,
+    the type or the targeted onnx opset.
+
+    .. code-block:: python
+
+        def qwen_version_selector(opset: int, *args: torch.Tensor) -> Tuple[str, torch.dtype]:
+            first_tensor = next(a for a in args if a is not None)
+            dtype = first_tensor.dtype
+            itype = torch_dtype_to_onnx_dtype(dtype)
+            if dtype == torch.float32:
+                if opset >= 23:
+                    return "LOOPA23", itype
+                return "LOOPMHA", itype
+            if dtype == torch.float16:
+                if first_tensor.is_cuda:
+                    return "PACKED", itype
+                return "LOOPMHA", itype
+            raise AssertionError(
+                f"Unable to handle type {torch.dtype} (itype={itype}) "
+                f"on device {torch.device} with opset={opset}"
+            )
+
+        qwen_sdpa_attention_versatile = EagerDirectReplacementWithOnnx(
+            qwen_sdpa_attention,
+            lambda qs, *args, **kwargs: torch.empty(
+                (qs.shape[0], qs.shape[2], qs.shape[1], qs.shape[3]),
+                dtype=qs.dtype,
+                device=qs.device,
+            ),
+            {
+                ("PACKED", onnx.TensorProto.FLOAT16): _add_com_microsoft_opset(
+                    PackedAttention.to_function_proto()
+                ),
+                ("LOOPA23", onnx.TensorProto.FLOAT): LoopAttention23.to_function_proto(),
+                ("LOOPA23", onnx.TensorProto.FLOAT16): _update_sequence_type(
+                    onnx.TensorProto.FLOAT16, LoopAttention23.to_function_proto()
+                ),
+                ("LOOPMHA", onnx.TensorProto.FLOAT): _add_com_microsoft_opset(
+                    LoopMHAAttention.to_function_proto()
+                ),
+                ("LOOPMHA", onnx.TensorProto.FLOAT16): _update_sequence_type(
+                    onnx.TensorProto.FLOAT16,
+                    _add_com_microsoft_opset(LoopMHAAttention.to_function_proto()),
+                ),
+            },
+            n_inputs=4,
+            n_outputs=1,
+            kwargs=dict(scaling=0.11180339887498948, num_heads=16),
+            name="qwen_sdpa_attention_versatile",
+            version_selector=qwen_version_selector,
+        )
     """
 
     def __init__(
         self,
         eager_fn: Callable[[TUPLE_TENSORS], TUPLE_TENSORS],
         shape_fn: Callable[[TUPLE_TENSORS], TUPLE_TENSORS],
-        function_proto: onnx.FunctionProto,
+        function_proto: Union[onnx.FunctionProto, Dict[Any, onnx.FunctionProto]],
         n_inputs: Optional[int] = None,
         n_outputs: Optional[int] = None,
         name: Optional[str] = None,
         kwargs: Optional[Dict[str, Union[int, float]]] = None,
         verbose: int = 0,
+        version_selector: Optional[Callable[..., Tuple[Any, ...]]] = None,
+        default_opset: int = 22,
     ):
-        assert isinstance(
-            function_proto, onnx.FunctionProto
+        assert isinstance(function_proto, onnx.FunctionProto) or (
+            isinstance(function_proto, dict)
+            or all(isinstance(v, onnx.FunctionProto) for v in function_proto.values())
         ), f"Unexpected type {type(function_proto)} for function_proto"
         assert isinstance(n_inputs, int), f"not implemented yet when n_inputs={n_inputs}"
-        assert isinstance(n_outputs, int), f"not implemented yet when n_inputs={n_outputs}"
+        assert isinstance(n_outputs, int), f"not implemented yet when n_outputs={n_outputs}"
         self.eager_fn = eager_fn
         self.shape_fn = shape_fn
-        self.function_proto = function_proto
+        self._function_proto = (
+            function_proto if isinstance(function_proto, onnx.FunctionProto) else None
+        )
+        self._function_proto_versioned = (
+            function_proto if isinstance(function_proto, dict) else {}
+        )
         self.n_inputs = n_inputs
         self.n_outputs = n_outputs
         self.name = name or (
@@ -170,24 +292,73 @@ class EagerDirectReplacementWithOnnx:
         )
         sig = inspect.signature(self.eager_fn)
         params = list(sig.parameters)
-        assert (
-            len(params) >= n_inputs
-        ), f"{self.eager_fn} accepts {params} as parameters < n_inputs={n_inputs}"
-        assert n_inputs == len(function_proto.input), (
-            f"Input mismatch n_inputs={n_inputs} but "
-            f"function_proto.input={function_proto.input}"
-        )
-        assert n_outputs == len(function_proto.output), (
-            f"Output mismatch n_outputs={n_outputs} but "
-            f"function_proto.output={function_proto.output}"
-        )
-        assert (
-            function_proto.domain == self.domain
-        ), f"Function domain must be {self.domain!r} but it is {function_proto.domain!r}"
         self.args_name = [p for p in params if p not in self.kwargs]
         self.kwargs_name = [p for p in params if p in self.kwargs]
         self.verbose = verbose
         self.custom_op = self._register()
+        self.version_selector = version_selector
+        self.default_opset = default_opset
+        self._check_protos(params)
+
+    def _check_protos(self, params):
+        assert (
+            len(params) >= self.n_inputs
+        ), f"{self.eager_fn} accepts {params} as parameters < n_inputs={self.n_inputs}"
+
+        # one proto
+        assert self._function_proto is None or self.n_inputs == len(
+            self._function_proto.input
+        ), (
+            f"Input mismatch n_inputs={self.n_inputs} but "
+            f"function_proto.input={self._function_proto.input}"
+        )
+        assert self._function_proto is None or self.n_outputs == len(
+            self._function_proto.output
+        ), (
+            f"Output mismatch n_outputs={self.n_outputs} but "
+            f"function_proto.output={self._function_proto.output}"
+        )
+        assert self._function_proto is None or (
+            self._function_proto.domain == self.domain
+        ), f"Function domain must be {self.domain!r} but it is {self._function_proto.domain!r}"
+
+        # multiple protos
+        assert all(
+            self.n_inputs == len(v.input) for v in self._function_proto_versioned.values()
+        ), f"Output mismatch n_inputs={self.n_inputs} but one version is wrong"
+        assert all(
+            self.n_outputs == len(v.output) for v in self._function_proto_versioned.values()
+        ), f"Output mismatch n_outputs={self.n_outputs} but one version is wrong"
+        assert all(
+            v.domain == self.domain for v in self._function_proto_versioned.values()
+        ), f"Function domain must be {self.domain!r} but it is different in one version"
+        assert (
+            not self._function_proto_versioned or self.version_selector
+        ), "version_selector is needed when multiple protos are given."
+
+    def get_function_proto(self, opset: int, *args) -> onnx.FunctionProto:
+        """Returns the correct version based on the inputs."""
+        if self._function_proto:
+            return self._function_proto
+        assert isinstance(
+            opset, int
+        ), f"The first argument must be an integer for the onnx opset but it is {type(opset)}"
+        assert any(
+            a is not None for a in args
+        ), f"Unexpected args={string_type(args, with_shape=True)}"
+        try:
+            key = self.version_selector(opset, *args)  # type: ignore[misc]
+        except (ValueError, AttributeError) as e:
+            raise AssertionError(
+                f"Unable to select a version, fails to get a key, available="
+                f"{set(self._function_proto_versioned)}, "
+                f"args={string_type(args,with_shape=True)}"
+            ) from e
+        assert key in self._function_proto_versioned, (
+            f"Unable to select a version, key={key}, available="
+            f"{set(self._function_proto_versioned)}, args={string_type(args,with_shape=True)}"
+        )
+        return self._function_proto_versioned[key]
 
     @property
     def domain(self) -> str:
@@ -219,6 +390,8 @@ class EagerDirectReplacementWithOnnx:
                 input_args.append(f"int {p}={val}")
             elif isinstance(val, float):
                 input_args.append(f"float {p}={val}")
+            elif isinstance(val, str):
+                input_args.append(f"str {p}={val}")
             else:
                 raise NotImplementedError(
                     f"kwargs {p!r} has a default value of unsupported type {type(val)}"
@@ -243,6 +416,7 @@ class EagerDirectReplacementWithOnnx:
         *args,
         engine: Optional[Callable] = None,
         dump_onnx_model: Optional[str] = None,
+        opset: int = 22,
         **kwargs,
     ) -> VerifyResult:
         """
@@ -257,6 +431,7 @@ class EagerDirectReplacementWithOnnx:
             :class:`onnx_diagnostic.reference.OnnxruntimeEvaluator`.
         :param dump_onnx_model: to dump the onnx model used to verify
             eager and onnx produce the same results
+        :param opset: onnx opset to use
         :param kwargs: additional arguments to the function
         :return: outputs of :func:`onnx_diagnostic.helpers.max_diff`
         """
@@ -291,7 +466,7 @@ class EagerDirectReplacementWithOnnx:
         assert engine is None, f"Not implemented yet with engine={engine!r}"
         ags, kws = self._make_args_kwargs(*args, **kwargs)
         sess = OnnxruntimeEvaluator(
-            self.function_proto,
+            self.get_function_proto(opset, *args),
             whole=True,
             dump_onnx_model=dump_onnx_model,
             function_kwargs=kws,
@@ -324,16 +499,25 @@ class EagerDirectReplacementWithOnnx:
             *args,
             **kwargs,
         ) -> Any:
-            if not g.has_local_function(
-                self.function_proto.name, domain=self.function_proto.domain
-            ):
-                g.add_function(self.function_proto)
+            has_devices = [a for a in args if isinstance(a, str) and g.has_device(a)]
+            assert (
+                has_devices
+            ), f"Missing device for any of the inputs {args}{g.get_debug_msg()}"
+            arg_device = has_devices[0]
+            fake_tensor = torch.empty(
+                tuple([(_ if isinstance(_, int) else 2) for _ in g.get_shape(args[0])]),
+                dtype=onnx_dtype_to_torch_dtype(g.get_type(args[0])),
+                device=int_device_to_torch_device(g.get_device(arg_device)),
+            )
+            function_proto = self.get_function_proto(g.main_opset, fake_tensor)
+            if not g.has_local_function(function_proto.name, domain=function_proto.domain):
+                g.add_function(function_proto)
             ags, kws = self._make_args_kwargs(*args, **kwargs)
             res = g.make_node(
-                self.function_proto.name,
+                function_proto.name,
                 ags,
                 outputs,
-                domain=self.function_proto.domain,
+                domain=function_proto.domain,
                 name=self.target_name,
                 **kws,
             )
@@ -356,41 +540,46 @@ class EagerDirectReplacementWithOnnx:
         """
         import onnxscript
 
-        onnx_plug_op = onnxscript.values.Opset(domain=self.function_proto.domain, version=1)
-        schema = onnx_plug_op[self.function_proto.name]
-        if schema is None:
-            all_types = [
-                "tensor(float)",
-                "tensor(float16)",
-                "tensor(bfloat16)",
-                "tensor(double)",
-                "tensor(int64)",
-                "tensor(int32)",
-            ]
-            type_constraints = []
-            for i in range(self.n_inputs):
-                type_constraints.append((f"T{i}", all_types, ""))
-            for i in range(self.n_outputs):
-                type_constraints.append((f"U{i}", all_types, ""))
-            schema = onnx.defs.OpSchema(
-                self.function_proto.name,
-                self.function_proto.domain,
-                1,
-                inputs=[
-                    onnx.defs.OpSchema.FormalParameter(f"arg_{i}", f"T{i}")
-                    for i in range(self.n_inputs)
-                ],
-                outputs=[
-                    onnx.defs.OpSchema.FormalParameter(f"res_{i}", f"U{i}")
-                    for i in range(self.n_outputs)
-                ],
-                type_constraints=type_constraints,
-            )
-            onnx.defs.register_schema(schema)
-        op = onnxscript.values.Op(onnx_plug_op, self.function_proto.name, schema)
+        onnx_plug_op = onnxscript.values.Opset(domain=self.domain, version=1)
+
+        def get_proto(*args):
+            function_proto = self.get_function_proto(self.default_opset, *args)
+            schema = onnx_plug_op[function_proto.name]
+            if schema is None:
+                all_types = [
+                    "tensor(float)",
+                    "tensor(float16)",
+                    "tensor(bfloat16)",
+                    "tensor(double)",
+                    "tensor(int64)",
+                    "tensor(int32)",
+                ]
+                type_constraints = []
+                for i in range(self.n_inputs):
+                    type_constraints.append((f"T{i}", all_types, ""))
+                for i in range(self.n_outputs):
+                    type_constraints.append((f"U{i}", all_types, ""))
+                schema = onnx.defs.OpSchema(
+                    function_proto.name,
+                    function_proto.domain,
+                    1,
+                    inputs=[
+                        onnx.defs.OpSchema.FormalParameter(f"arg_{i}", f"T{i}")
+                        for i in range(self.n_inputs)
+                    ],
+                    outputs=[
+                        onnx.defs.OpSchema.FormalParameter(f"res_{i}", f"U{i}")
+                        for i in range(self.n_outputs)
+                    ],
+                    type_constraints=type_constraints,
+                )
+                onnx.defs.register_schema(schema)
+            op = onnxscript.values.Op(onnx_plug_op, function_proto.name, schema)
+            return op
 
         def converter(*cargs, **ckwargs):
             ags, kws = self._make_args_kwargs(*cargs, **ckwargs)
+            op = get_proto(*cargs)
             return op(*ags, n_outputs=self.n_outputs, **kws)
 
         return onnxscript.values.TracedOnnxFunction(onnx_plug_op, converter)

onnx_diagnostic/ext_test_case.py

@@ -610,6 +610,21 @@ def requires_onnxruntime(version: str, msg: str = "") -> Callable:
     return lambda x: x
 
 
+def has_onnxruntime(version: str, msg: str = "") -> Callable:
+    """Skips a unit test if :epkg:`onnxruntime` is not recent enough."""
+    import packaging.version as pv
+    import onnxruntime
+
+    if not hasattr(onnxruntime, "__version__"):
+        # development version
+        return True
+
+    if pv.Version(onnxruntime.__version__) < pv.Version(version):
+        msg = f"onnxruntime version {onnxruntime.__version__} < {version}: {msg}"
+        return False
+    return True
+
+
 def has_onnxruntime_training(push_back_batch: bool = False):
     """Tells if onnxruntime_training is installed."""
     try:
@@ -830,6 +845,13 @@ class ExtTestCase(unittest.TestCase):
             f.write(proto.SerializeToString())
         return fullname
 
+    def dump_text(self, name: str, text: str, folder: Optional[str] = None) -> str:
+        """Dumps text in a file."""
+        fullname = self.get_dump_file(name, folder=folder)
+        with open(fullname, "w") as f:
+            f.write(text)
+        return fullname
+
     def assertExists(self, name):
         """Checks the existing of a file."""
         if not os.path.exists(name):
@@ -1196,9 +1218,9 @@ class ExtTestCase(unittest.TestCase):
     def assert_onnx_disc(
         self,
         test_name: str,
-        proto: "onnx.ModelProto",  # noqa: F821
+        proto: Union[str, "onnx.ModelProto"],  # noqa: F821
         model: "torch.nn.Module",  # noqa: F821
-        inputs: Union[Tuple[Any], Dict[str, Any]],
+        inputs: Union[Tuple[Any], Dict[str, Any], List[Any]],
         verbose: int = 0,
         atol: float = 1e-5,
         rtol: float = 1e-3,
@@ -1242,7 +1264,9 @@ class ExtTestCase(unittest.TestCase):
             name = f"{test_name}.onnx"
             if verbose:
                 print(f"[{vname}] save the onnx model into {name!r}")
+            model_file = None
             if isinstance(proto, str):
+                model_file = proto
                 name = proto
                 proto = onnx.load(name)
             elif not self.unit_test_going():
@@ -1255,45 +1279,64 @@ class ExtTestCase(unittest.TestCase):
         if verbose:
             print(f"[{vname}] make feeds {string_type(inputs, **kws)}")
 
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+            if expected is not None:
+                expected = [expected]
+
+        gots = []
         if use_ort:
             assert isinstance(
                 proto, onnx.ModelProto
             ), f"Unexpected type {type(proto)} for proto"
-            feeds = make_feeds(proto, inputs, use_numpy=True, copy=True)
             import onnxruntime
 
             options = onnxruntime.SessionOptions()
             if ort_optimized_graph:
                 options.optimized_model_filepath = f"{name}.optort.onnx"
+            if "log_severity_level" in kwargs:
+                options.log_severity_level = kwargs["log_severity_level"]
+            if "log_verbosity_level" in kwargs:
+                options.log_verbosity_level = kwargs["log_verbosity_level"]
             providers = kwargs.get("providers", ["CPUExecutionProvider"])
             if verbose:
                 print(f"[{vname}] create onnxruntime.InferenceSession with {providers}")
             sess = onnxruntime.InferenceSession(
-                proto.SerializeToString(), options, providers=providers
+                model_file or proto.SerializeToString(), options, providers=providers
             )
-            if verbose:
-                print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")
-            got = sess.run(None, feeds)
+            for inp in inputs:
+                feeds = make_feeds(proto, inp, use_numpy=True, copy=True)
+                if verbose:
+                    print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")
+                got = sess.run(None, feeds)
+                gots.append(got)
         else:
-            feeds = make_feeds(proto, inputs, copy=True)
             if verbose:
                 print(f"[{vname}] create InferenceSessionForTorch")
             sess = InferenceSessionForTorch(proto, **kwargs)
-            if verbose:
-                print(f"[{vname}] run orttorch feeds {string_type(feeds, **kws)}")
-            got = sess.run(None, feeds)
+            for inp in inputs:
+                feeds = make_feeds(proto, inp, copy=True)
+                if verbose:
+                    print(f"[{vname}] run orttorch feeds {string_type(feeds, **kws)}")
+                got = sess.run(None, feeds)
+                gots.append(got)
         if verbose:
             print(f"[{vname}] compute expected values")
 
         if expected is None:
             if copy_inputs:
-                expected = (
-                    model(*copy.deepcopy(inputs))
-                    if isinstance(inputs, tuple)
-                    else model(**copy.deepcopy(inputs))
-                )
+                expected = [
+                    (
+                        model(*copy.deepcopy(inp))
+                        if isinstance(inp, tuple)
+                        else model(**copy.deepcopy(inp))
+                    )
+                    for inp in inputs
+                ]
             else:
-                expected = model(*inputs) if isinstance(inputs, tuple) else model(**inputs)
+                expected = [
+                    model(*inp) if isinstance(inp, tuple) else model(**inp) for inp in inputs
+                ]
 
         if verbose:
             print(f"[{vname}] expected {string_type(expected, **kws)}")
@@ -1306,47 +1349,50 @@ class ExtTestCase(unittest.TestCase):
                 import torch
 
                 ep = torch.export.load(ep)
-            ep_inputs = copy.deepcopy(inputs) if copy_inputs else inputs
+
             ep_model = ep.module()  # type: ignore[union-attr]
-            ep_expected = (
-                ep_model(*copy.deepcopy(ep_inputs))
-                if isinstance(ep_inputs, tuple)
-                else ep_model(**copy.deepcopy(ep_inputs))
-            )
-            if verbose:
-                print(f"[{vname}] ep_expected {string_type(ep_expected, **kws)}")
-            ep_diff = max_diff(expected, ep_expected, hist=[0.1, 0.01])
+            for expe, inp, got in zip(expected, inputs, gots):
+                ep_inputs = copy.deepcopy(inp) if copy_inputs else inp
+                ep_expected = (
+                    ep_model(*copy.deepcopy(ep_inputs))
+                    if isinstance(ep_inputs, tuple)
+                    else ep_model(**copy.deepcopy(ep_inputs))
+                )
+                if verbose:
+                    print(f"[{vname}] ep_expected {string_type(ep_expected, **kws)}")
+                ep_diff = max_diff(expe, ep_expected, hist=[0.1, 0.01])
+                if verbose:
+                    print(f"[{vname}] ep_diff {string_diff(ep_diff)}")
+                assert (
+                    isinstance(ep_diff["abs"], float)
+                    and isinstance(ep_diff["rel"], float)
+                    and not numpy.isnan(ep_diff["abs"])
+                    and ep_diff["abs"] <= atol
+                    and not numpy.isnan(ep_diff["rel"])
+                    and ep_diff["rel"] <= rtol
+                ), (
+                    f"discrepancies in {test_name!r} between the exported program "
+                    f"and the exported model diff={string_diff(ep_diff)}"
+                )
+                ep_nx_diff = max_diff(ep_expected, got, flatten=True, hist=[0.1, 0.01])
+                if verbose:
+                    print(f"[{vname}] ep_nx_diff {string_diff(ep_nx_diff)}")
+
+        for expe, got in zip(expected, gots):
+            diff = max_diff(expe, got, flatten=True, hist=[0.1, 0.01])
             if verbose:
-                print(f"[{vname}] ep_diff {string_diff(ep_diff)}")
+                print(f"[{vname}] diff {string_diff(diff)}")
             assert (
-                isinstance(ep_diff["abs"], float)
-                and isinstance(ep_diff["rel"], float)
-                and not numpy.isnan(ep_diff["abs"])
-                and ep_diff["abs"] <= atol
-                and not numpy.isnan(ep_diff["rel"])
-                and ep_diff["rel"] <= rtol
+                isinstance(diff["abs"], float)
+                and isinstance(diff["rel"], float)
+                and not numpy.isnan(diff["abs"])
+                and diff["abs"] <= atol
+                and not numpy.isnan(diff["rel"])
+                and diff["rel"] <= rtol
             ), (
-                f"discrepancies in {test_name!r} between the exported program "
-                f"and the exported model diff={string_diff(ep_diff)}"
+                f"discrepancies in {test_name!r} between the model and "
+                f"the onnx model diff={string_diff(diff)}"
             )
-            ep_nx_diff = max_diff(ep_expected, got, flatten=True, hist=[0.1, 0.01])
-            if verbose:
-                print(f"[{vname}] ep_nx_diff {string_diff(ep_nx_diff)}")
-
-        diff = max_diff(expected, got, flatten=True, hist=[0.1, 0.01])
-        if verbose:
-            print(f"[{vname}] diff {string_diff(diff)}")
-        assert (
-            isinstance(diff["abs"], float)
-            and isinstance(diff["rel"], float)
-            and not numpy.isnan(diff["abs"])
-            and diff["abs"] <= atol
-            and not numpy.isnan(diff["rel"])
-            and diff["rel"] <= rtol
-        ), (
-            f"discrepancies in {test_name!r} between the model and "
-            f"the onnx model diff={string_diff(diff)}"
-        )
 
     def _debug(self):
         "Tells if DEBUG=1 is set up."