add argmax

Author: xadupre

onnx_diagnostic/helpers/helper.py

@@ -1181,6 +1181,7 @@ def max_diff(
                 if exp_cpu.size == got_cpu.size
                 else (np.inf, np.inf, np.inf, 0, np.inf)
             )
+            argm = None
         else:
             abs_diff, rel_diff, sum_diff, n_diff, nan_diff = (
                 float(diff.max()),
@@ -1189,6 +1190,7 @@ def max_diff(
                 float(diff.size),
                 float(ndiff.sum()),
             )
+            argm = tuple(map(int, np.unravel_index(diff.argmax(), diff.shape)))
         if verbose >= 10 and (abs_diff >= 10 or rel_diff >= 10):
             # To understand the value it comes from.
             if debug_info:
@@ -1219,7 +1221,9 @@ def max_diff(
                     f"_index={_index}"
                 )
 
-        res = dict(abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff)
+        res = dict(
+            abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
+        )
         if hist:
             if isinstance(hist, bool):
                 hist = np.array([0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype)
@@ -1284,8 +1288,10 @@ def max_diff(
                 float(diff.numel()),
                 float(ndiff.sum()),
             )
+            argm = tuple(map(int, torch.unravel_index(diff.argmax(), diff.shape)))
         elif got_cpu.numel() == exp_cpu.numel():
             abs_diff, rel_diff, sum_diff, n_diff, nan_diff = (0.0, 0.0, 0.0, 0.0, 0.0)
+            argm = None
         else:
             abs_diff, rel_diff, sum_diff, n_diff, nan_diff = (
                 np.inf,
@@ -1294,6 +1300,7 @@ def max_diff(
                 np.inf,
                 np.inf,
             )
+            argm = None
 
         if verbose >= 10 and (abs_diff >= 10 or rel_diff >= 10):
             # To understand the value it comes from.
@@ -1325,7 +1332,9 @@ def max_diff(
                     f"_index={_index}"
                 )
 
-        res = dict(abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff)
+        res = dict(
+            abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
+        )
         if hist:
             if isinstance(hist, bool):
                 hist = torch.tensor(
@@ -1478,6 +1487,8 @@ def string_diff(diff: Dict[str, Any]) -> str:
                 rows.append(f"#{v}{k}")
         suffix = "-".join(rows)
         suffix = f"/{suffix}"
+    if "argm" in diff:
+        suffix += f", argmax={diff['argm']}"
     if diff.get("dnan", None):
         if diff["abs"] == 0 or diff["rel"] == 0:
             return f"abs={diff['abs']}, rel={diff['rel']}, dnan={diff['dnan']}{suffix}"