Fix evaluating colormaps on non-numpy arrays

timhoffm · timhoffm · commit be849c5680b3 · 2022-09-25T22:34:17.000+02:00
Closes #23132. I can't specifically test that case, because we don't have pytorch as a test dependency. However, I'd claim that deferring np.nan() to after converting to a numpy array is obviously the right thing to do.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -706,8 +706,12 @@ def __call__(self, X, alpha=None, bytes=False):
         if not self._isinit:
             self._init()
 
-        mask_bad = X.mask if np.ma.is_masked(X) else np.isnan(X)  # Mask nan's.
+        # Take the bad mask from a masked array, or in all other cases defer
+        # np.isnan() to after we have converted to an array.
+        mask_bad = X.mask if np.ma.is_masked(X) else None
         xa = np.array(X, copy=True)
+        if mask_bad is None:
+            mask_bad = np.isnan(xa)
         if not xa.dtype.isnative:
             xa = xa.byteswap().newbyteorder()  # Native byteorder is faster.
         if xa.dtype.kind == "f":
