bfloat16 arrays are incompatible with pickle

[jakevdp]

import pickle
import jax.numpy as jnp
pickle.dumps(jnp.bfloat16(0))

---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
<ipython-input-13-9fa7e9015eb7> in <module>()
      1 import pickle
      2 import jax.numpy as jnp
----> 3 pickle.dumps(jnp.bfloat16(0))

PicklingError: Can't pickle <class 'bfloat16'>: attribute lookup bfloat16 on builtins failed

I think this is due to the fact that picle the bfloat16 custom type does not correctly define its full class name. Compare these::

from jax._src.lib import xla_client
print(type(xla_client.bfloat16(0.0)))
# <class 'bfloat16'>

print(type(jnp.array(0)))
# <class 'jaxlib.xla_extension.DeviceArray'>

print(type(np.float16(0))
# <class 'numpy.float16'>

[jakevdp]

I think for this to work correctly, we'd need the tp_name attribute here to correspond to its actual import location within JAX, i.e. jaxlib.xla_client.bfloat16.

That's not straightforward, though, because this type definition is shared with tensorflow and exported at a different package path within that library.

[jakevdp]

Indeed the tensorflow bfloat16 dtype has the same issue:

import tensorflow as tf
import pickle
import numpy as np
pickle.dumps(np.ones(10, dtype=tf.bfloat16.as_numpy_dtype))

---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
<ipython-input-24-272a81341603> in <module>()
      2 import pickle
      3 import numpy as np
----> 4 pickle.dumps(np.ones(10, dtype=tf.bfloat16.as_numpy_dtype))

PicklingError: Can't pickle <class 'bfloat16'>: attribute lookup bfloat16 on builtins failed

[jakevdp]

I think once this bfloat16 name issue is resolved, numpy's pickle serialization/deserialization should work. For example:

import jax.numpy as jnp
x = jnp.arange(4, dtype='bfloat16')
constructor, args, state = x.__reduce__()
x2 = constructor(*args)
x2.__setstate__(state)
print(repr(x2))
# array([0, 1, 2, 3], dtype=bfloat16)

The error here comes from when pickle tries to serialize the bfloat16 dtype in state, for the reasons described above.

[jakevdp]

Here's a hack that can be used to work around this:

import builtins
import pickle
from jax._src.lib import xla_client
import jax.numpy as jnp

# Hack: this is the module reported by this object.
builtins.bfloat16 = xla_client.bfloat16

s = pickle.dumps(jnp.arange(10, dtype='bfloat16'))
y = pickle.loads(s)
print(repr(y))
# array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=bfloat16)

One thing we could do to fix this (though it's kind of terrible) is to set this attribute of builtins when jax is imported.

[hamzamerzic]

A marginally simpler hack that we've been using is

builtins.bfloat16 = jnp.dtype('bfloat16').type
instead of
builtins.bfloat16 = xla_client.bfloat16,

as it avoids the private from._src.lib import.

Setting builtins.bfloat16 = jnp.bfloat16 does not seem to work, even though xla_client.bfloat16 == jnp.bfloat16 evaluates to True.

[hawkinsp]

#15122 fixed this!

The type is now just ml_dtypes.bfloat16, and it should be pickleable (there is a test for it).

[jakevdp]

Pickle is solved, but there are still problems with numpy-native array serialization: jax-ml/ml_dtypes#41