Buffer protocol for jaxlib.xla_extension.Array

[heiner]

Googling suggests a Python version of PyObject_CheckBuffer could look like

def check_buffer_doesntwork(obj):
    try:
        memoryview(obj)
    except TypeError:
        return False
    return True

This works for np.ndarrays and various other types, but not for jaxlib.xla_extension.Array:

>>> memoryview(jnp.arange(10))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: memoryview: a bytes-like object is required, not 'Array'

Is this intentional?

[hawkinsp]

Hmm. We used to support this. I think it might have gotten lost as part of the unified jax.Array transition in 0.4. Let me look into it.

[hawkinsp]

Created #14713 to track.

[heiner]

Thanks! Great turnaround time! :)