Add support for checking np.floating as well as float when summing at…

…tributions. PiperOrigin-RevId: 405737562
tensorflow · Oct 26, 2021 · 80a7dfb · 80a7dfb
1 parent 0c66788
commit 80a7dfb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tensorflow_model_analysis/metrics/attributions.py b/tensorflow_model_analysis/metrics/attributions.py
@@ -276,7 +276,8 @@ def __init__(self, key: metric_types.AttributionsKey, absolute: bool):
 
   def _sum(self, a: List[float], b: Union[np.ndarray, List[float]]):
     """Adds values in b to a at matching offsets."""
-    if isinstance(b, float) or (isinstance(b, np.ndarray) and b.size == 1):
+    if (isinstance(b, (float, np.floating)) or
+        (isinstance(b, np.ndarray) and b.size == 1)):
       if len(a) != 1:
         raise ValueError(
             'Attributions have different array sizes {} != {}'.format(a, b))