Compatibility of `DeviceArray` with `PIL.Image.fromarray()` to mirror numpy behavior

This is just a “quality of life” feature request, but twice now I’ve had to import the real numpy to do this sort of thing:

PIL.Image.fromarray(np.array(jnp_array)).show()

It would be nice to be able to avoid importing the real numpy and just write:

PIL.Image.fromarray(jnp_array).show()

Apologies if this is actually trivially achievable without importing the real numpy - I’m new to Python. Seems possible since no one has raised this issue yet. But then again, it’d be nice to have compatibility if it’s possible, since this Image.fromarray(numpy_array).show() pattern seems common?

Thanks!

[hawkinsp]

JAX doesn’t implement __array_interface__, but it does implement the PEP 3118 buffer protocol which is extremely similar. We could support both, most likely.

[shoyer]

I see no reason why PIL should not handle __array__. This would also give them out of the box support for other array-convertible types in Python like those from Xarray and Dask, which seems nice.

On the other hand, I guess JAX could also easily implement __array_interface__, likely by converting into a NumPy array and then calling ndarray. __array_interface__().