Repository: matplotlib/matplotlib
Issue category: Performance Issue

Issue:
[Bug]: plt.hist takes significantly more time with torch and jax arrays
### Bug summary

Hi,

Time taken to plot `plt.hist` directly on `jax` or `torch` arrays is significantly more than **combined time taken to first convert them to `numpy` and then using `plt.hist`**. Shouldn't `matplotlib` internally convert them to `numpy` arrays before plotting?

To reproduce the bug, directly run the following snippet on Google Colab.

### Code for reproduction

```python
from time import time
import numpy as np

import torch

import jax
import jax.random as jr
import jax.numpy as jnp

import matplotlib.pyplot as plt

jax_array = jr.normal(jr.PRNGKey(0), (1000, 150))
torch_array = torch.randn(1000, 150)

def plot_hist(array):
    init = time()
    plt.figure()
    plt.hist(array)
    print(f"Time to plot: {time() - init:.2f} s")
    plt.show()
    
plot_hist(jax_array.ravel())
plot_hist(torch_array.ravel())
plot_hist(np.array(jax_array.ravel()))
plot_hist(np.array(torch_array.ravel()))
```


### Actual outcome

Time to plot: 4.19 s
![image](https://github.com/matplotlib/matplotlib/assets/59758528/29f7e6e6-4c91-4801-8472-24eaaae591cc)

Time to plot: 2.61 s
![image](https://github.com/matplotlib/matplotlib/assets/59758528/cf0b12c2-fabf-45fb-8e59-89206bb6ed68)

Time to plot: 0.03 s
![image](https://github.com/matplotlib/matplotlib/assets/59758528/9a39702e-b0ba-48ba-9ec6-5731f3d4b372)

Time to plot: 0.04 s
![image](https://github.com/matplotlib/matplotlib/assets/59758528/5e99fa5c-c614-4422-9b98-ff62fa3d3563)


### Expected outcome

Time to plot: 0.03 s

Time to plot: 0.04 s

Time to plot: 0.03 s

Time to plot: 0.04 s

### Additional information

> What are the conditions under which this bug happens? input parameters, edge cases, etc?

It is happening with all kinds of shapes.

> Has this worked in earlier versions?

Tested with default colab `matplotlib` version 3.7.1 and also with 3.6.3.

> Do you know why this bug is happening?

Not exactly sure.

> Do you maybe even know a fix?

Maybe convert any python object to a `numpy` array before plotting?

### Operating system

Ubuntu 20.04.5 LTS

### Matplotlib Version

3.7.1

### Matplotlib Backend

module://matplotlib_inline.backend_inline

### Python version

3.10.11

### Jupyter version

6.4.8

### Installation

None


Hints:
The unpacking happens here:
https://github.com/matplotlib/matplotlib/blob/b61bb0b6392c23d38cd45c658bfcd44df145830d/lib/matplotlib/cbook.py#L2237-L2251

The pytorch tensor does not support any of the conversion methods, so Matplotlib doesn't really know what to do with it. There is a discussion in https://github.com/matplotlib/matplotlib/issues/22645 about this, but if I remember correctly we expect the libraries to support the `to_numpy` method (but still support the `values` attribute).

(I could not install jax, but I suppose something similar goes on there.)
And when the conversion doesn't work, it ends up in this loop:
https://github.com/matplotlib/matplotlib/blob/b61bb0b6392c23d38cd45c658bfcd44df145830d/lib/matplotlib/cbook.py#L1332-L1348
which is where most of the time is spent.
Thanks for the quick response, @oscargus! Given that both these libraries support `.__array__()` method for conversion to numpy array, wouldn't it be easier to add one more `if` condition in `_unpack_to_numpy` to include them?
```py
type(jax_array.__array__()), type(torch_array.__array__())
# Output: (numpy.ndarray, numpy.ndarray)
```
Yes, I also noted that. It probably can make sense.

(I think the reason why we do this somewhat carefully is for unit information to not get lost.)

Would you be interested in submitting a patch? I think that if this goes last in the conversion chain, it shouldn't break too many things... (A problem here is that we do not, yet, test for "all" types that possibly can be used and "works". There's been a discussion of having a special test suite for that, but it has not yet been implemented.)
Even `tensorflow` supports `__array__()` method. I guess these 3 libraries account for almost 99% of the machine learning codebase available online :) It'd be great if this conversion passes without breaking many things!

Sure, I'd submit a patch. I guess I need to only change the `_unpack_to_numpy` to the following, right?

```py
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 
         return x 
     if hasattr(x, 'to_numpy'): 
         # Assume that any to_numpy() method actually returns a numpy array 
         return x.to_numpy() 
     if hasattr(x, 'values'): 
         xtmp = x.values 
         # For example a dict has a 'values' attribute, but it is not a property 
         # 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)
         return x.__array__()
     return x 
```
Yes, but please verify that `__array__` actually returns a numpy array, like we do with `values` above.
Thank you for the important suggestion, @timhoffm. Now, `__array__` method check works in theory for the cases I imagined but `np.float32` type objects get stuck into that check. When `__array__` method is called on `np.float32` object, it gets converted to `ndarray` type and eventually this leads to an infinite recursion. 

A temporary fix I could figure out is to add two more `if` conditions to check if object is of type `np.floating` (includes all float types) and type `np.integer` (includes all integer types including `uint`). I can also include a boolean check. Will it be all or this already looks unpythonic?

More directions to solve this issue could be the following:
1. Raise an issue to add `to_numpy()` methods in JAX and PyTorch repos.
2. Raise an issue to have a universal `numpy` object checker type in NumPy library so that we can replace `ndarray` check with that. After this, any numpy object will be captured in the first check. 
3. Add hard-coded checks for JAX and PyTorch like the following:

```py
if str(type(x)) == "<class 'torch.Tensor'>":
    return x.__array__()
if str(type(x)) == "<class 'jaxlib.xla_extension.ArrayImpl'>":
    return x.__array__()
```

I am open to your suggestions.

Edit1: [`np.generic`](https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.generic) works for most ([all?](https://numpy.org/doc/stable/reference/arrays.scalars.html#numpy.generic)) scalars, so we can add `if isinstance(x, np.generic)` as the second check just after `ndarray` check like the following:

```py
def _unpack_to_numpy(x):
    """Internal helper to extract data from e.g. pandas and xarray objects."""
    if isinstance(x, np.ndarray):
        # 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
        return x.to_numpy()
    if hasattr(x, 'values'):
        xtmp = x.values
        # For example a dict has a 'values' attribute, but it is not a property
        # 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
```

Patch:
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
index ff6b2a15ec35..9d03752c666c 100644
--- a/lib/matplotlib/cbook.py
+++ b/lib/matplotlib/cbook.py
@@ -2344,6 +2344,30 @@ def _picklable_class_constructor(mixin_class, fmt, attr_name, base_class):
     return cls.__new__(cls)
 
 
+def _is_torch_array(x):
+    """Check if 'x' is a PyTorch Tensor."""
+    try:
+        # we're intentionally not attempting to import torch. If somebody
+        # has created a torch array, torch should already be in sys.modules
+        return isinstance(x, sys.modules['torch'].Tensor)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
+
+
+def _is_jax_array(x):
+    """Check if 'x' is a JAX Array."""
+    try:
+        # we're intentionally not attempting to import jax. If somebody
+        # has created a jax array, jax should already be in sys.modules
+        return isinstance(x, sys.modules['jax'].Array)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
+
+
 def _unpack_to_numpy(x):
     """Internal helper to extract data from e.g. pandas and xarray objects."""
     if isinstance(x, np.ndarray):
@@ -2358,6 +2382,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 _is_torch_array(x) or _is_jax_array(x):
+        xtmp = x.__array__()
+
+        # In case __array__() method does not return a numpy array in future
+        if isinstance(xtmp, np.ndarray):
+            return xtmp
     return x
 
 


Test patch:
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
index 1f4f96324e9e..24fd02e65a5f 100644
--- a/lib/matplotlib/tests/test_cbook.py
+++ b/lib/matplotlib/tests/test_cbook.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import sys
 import itertools
 import pickle
 
@@ -16,6 +17,7 @@
 from matplotlib import _api, cbook
 import matplotlib.colors as mcolors
 from matplotlib.cbook import delete_masked_points, strip_math
+from types import ModuleType
 
 
 class Test_delete_masked_points:
@@ -938,3 +940,45 @@ def test_auto_format_str(fmt, value, result):
     """Apply *value* to the format string *fmt*."""
     assert cbook._auto_format_str(fmt, value) == result
     assert cbook._auto_format_str(fmt, np.float64(value)) == result
+
+
+def test_unpack_to_numpy_from_torch():
+    """Test that torch tensors are converted to numpy arrays.
+    We don't want to create a dependency on torch in the test suite, so we mock it.
+    """
+    class Tensor:
+        def __init__(self, data):
+            self.data = data
+        def __array__(self):
+            return self.data
+    torch = ModuleType('torch')
+    torch.Tensor = Tensor
+    sys.modules['torch'] = torch
+
+    data = np.arange(10)
+    torch_tensor = torch.Tensor(data)
+
+    result = cbook._unpack_to_numpy(torch_tensor)
+    assert result is torch_tensor.__array__()
+
+
+def test_unpack_to_numpy_from_jax():
+    """Test that jax arrays are converted to numpy arrays.
+    We don't want to create a dependency on jax in the test suite, so we mock it.
+    """
+    class Array:
+        def __init__(self, data):
+            self.data = data
+        def __array__(self):
+            return self.data
+
+    jax = ModuleType('jax')
+    jax.Array = Array
+
+    sys.modules['jax'] = jax
+
+    data = np.arange(10)
+    jax_array = jax.Array(data)
+
+    result = cbook._unpack_to_numpy(jax_array)
+    assert result is jax_array.__array__()