Repository: matplotlib/matplotlib
Issue category: Performance 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


Functions to edit:
lib/matplotlib/cbook.py:_unpack_to_numpy
Functions added: lib/matplotlib/cbook.py:_is_torch_array
lib/matplotlib/cbook.py:_is_jax_array

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__()