Merge pull request #24009 from timhoffm/fix-pytorch

tacaswell · web-flow · commit 30fbb45784b7 · 2022-10-20T17:51:04.000-04:00
Fix evaluating colormaps on non-numpy arrays
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":
