update _unpack_to_numpy to handle JAX and PyTorch arrays

Author: patel-zeel

lib/matplotlib/cbook.py

@@ -2237,7 +2237,10 @@ def _picklable_class_constructor(mixin_class, fmt, attr_name, base_class):
 def _unpack_to_numpy(x):
     """Internal helper to extract data from e.g. pandas and xarray objects."""
     if isinstance(x, np.ndarray):
-        # If numpy, return directly
+        # If numpy array, return directly
+        return x
+    if isinstance(x, np.generic):
+        # If numpy scalar, return directly
         return x
     if hasattr(x, 'to_numpy'):
         # Assume that any to_numpy() method actually returns a numpy array
@@ -2248,6 +2251,12 @@ def _unpack_to_numpy(x):
         # so in this case we do not want to return a function
         if isinstance(xtmp, np.ndarray):
             return xtmp
+    if hasattr(x, '__array__'):
+        # Assume that any to __array__() method returns a numpy array (e.g. TensorFlow, JAX or PyTorch arrays)
+        x = x.__array__()
+        # Anything that doesn't return ndarray via __array__() method will be filtered by the following check
+        if isinstance(x, np.ndarray):
+            return x
     return x