mypy

Author: xadupre

onnx_diagnostic/ext_test_case.py

@@ -1195,7 +1195,9 @@ def assert_onnx_disc(
         if verbose:
             print(f"[{vname}] diff {string_diff(diff)}")
         assert (
-            not numpy.isnan(diff["abs"])
+            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

onnx_diagnostic/helpers/helper.py

@@ -1221,7 +1221,7 @@ def max_diff(
                     f"_index={_index}"
                 )
 
-        res: Dict[str, Union[int, float, Tuple[int, ...]]] = dict(
+        res: Dict[str, float] = dict(
             abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
         )
         if hist:
@@ -1332,7 +1332,7 @@ def max_diff(
                     f"_index={_index}"
                 )
 
-        res: Dict[str, Union[int, float, Tuple[int, ...]]] = dict(
+        res: Dict[str, float] = dict(
             abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
         )
         if hist: