Monkeypatching matplotlib

[syamajala]

There is a lot of user code out there that uses matplotlib convenience functions which wrap numpy. It would be nice if we could monkeypatch matplotlib to use cunumeric.

Here is an example:

import cunumeric as np

import matplotlib.axes
matplotlib.axes._axes.np = np

import matplotlib.pyplot as plt
plt.np = np

from matplotlib.colors import LogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable

events = 100000
fig, ax = plt.subplots(ncols=3)
data = np.random.randn(events)
for i in range(0, 3):
    ax[i].hist(np.random.randn(events), bins=10)

plt.savefig("hist.png")

Here is the error I see when I do that:

(/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly) [seshu@sdfampere018 test]$ legate --gpus 4 --profile test.py
/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/_utils/coverage.py:177: RuntimeWarning: cuNumeric has not implemented numpy.iterable and is falling back to canonical NumPy. You may notice significantly decreased performance for this function call.
  warnings.warn(
Traceback (most recent call last):
  File "/sdf/home/s/seshu/dev/test/test.py", line 15, in <module>
    ax[i].hist(np.random.randn(events), bins=10)
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/matplotlib/__init__.py", line 1473, in inner
    return func(
           ^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/matplotlib/axes/_axes.py", line 7061, in hist
    bars = _barfunc(bins[:-1]+boffset, height, width,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/matplotlib/__init__.py", line 1473, in inner
    return func(
           ^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/matplotlib/axes/_axes.py", line 2520, in bar
    x, height, width, y, linewidth, hatch = np.broadcast_arrays(
                                            ^^^^^^^^^^^^^^^^^^^^
  File "runtime.pyx", line 1132, in legate._lib.runtime.runtime.track_provenance.decorator.wrapper
  File "runtime.pyx", line 1133, in legate._lib.runtime.runtime.track_provenance.decorator.wrapper
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/_utils/coverage.py", line 111, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/_module/array_dimension.py", line 317, in broadcast_arrays
    arrs = [convert_to_cunumeric_ndarray(arr) for arr in args]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/_array/util.py", line 130, in convert_to_cunumeric_ndarray
    thunk = runtime.get_numpy_thunk(obj, share=share)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/runtime.py", line
295, in get_numpy_thunk
    return self.find_or_create_array_thunk(obj, transfer)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sdf/group/lcls/ds/tools/conda_envs/cunumeric-nightly/lib/python3.12/site-packages/cunumeric/runtime.py", line
379, in find_or_create_array_thunk
    raise TypeError(f"cuNumeric does not support dtype={array.dtype}")
TypeError: cuNumeric does not support dtype=object

I know cunumeric does not support object arrays, but would it be possible to fall back to numpy in that case?

[lightsighter]

Rather than monkey-patching matplotlib, it feels like a better answer would be to use NumPy overrides.

I know cunumeric does not support object arrays, but would it be possible to fall back to numpy in that case?

We probably should be falling back in that case.

[manopapad]

Rather than monkey-patching matplotlib, it feels like a better answer would be to use NumPy overrides.

We have that today, it doesn't seem to work for this case, so matplotlib is probably doing something different than just

numpy.histogram(some_array_that_could_be_cunumeric_used_directly)

E.g. it's explicitly calling np.asarray before doing anything else. There is another issue open to do a full investigation of why this happening.

[syamajala]

I looked at the matplotlib code some more and its difficult to follow, but I think there are calls to np.asarray before it gets to the np.histogram within plt.hist().

[manopapad]

I don't know for sure if np.asarray is the culprit here, but @seberg would we be within our rights to return a cuNumeric array when np.asarray is called? I'm not even sure if np.asarray is one of the functions we can override through __array_function__.

Also, something that wasn't mentioned yet, we do have an lgpatch utility that will do this monkeypatching, but as you can see it won't necessarily be a seamless experience.

[seberg]

You can't return a cunumeric array from np.asarray() unless the user passed np.asarray(arr, like=cunumeric_array).
There is a bit of a missing np.asduckarray() for a library like matplotlib to use; even if that is trivial to implement and probably exists indirectly via other functions...