Issues with saving and pickling bfloat16 array

[ivandustin]

Hi,

Saving a bfloat16 array using jax.numpy.save seems to have a problem.

In [162]: a = jnp.array([1, 2, 3], dtype=jnp.bfloat16)
In [163]: a
Out[163]: DeviceArray([1, 2, 3], dtype=bfloat16)
In [164]: jnp.save('test.npy', a)
In [165]: b = jnp.load('test.npy')
In [166]: b
Out[166]: array([b'\x80\x3F', b'\x00\x40', b'\x40\x40'], dtype='|V2')

As you can see, the result of jax.numpy.load seems to be garbage.

There is also a problem when using pickle.

In [167]: import pickle
In [168]: pickle.dumps(a)
---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
<ipython-input-168-e346a7b9f67a> in <module>
----> 1 pickle.dumps(a)
PicklingError: Can't pickle <class 'bfloat16'>: attribute lookup bfloat16 on builtins failed

Is this a bug?

Thank you.

jaxlib version is 0.1.73

[jakevdp]

Thanks for the report - the issue is that bfloat16 is not a native numpy type. If you want to recover the original array, you can do so by viewing the output as type jnp.bfloat16:

import jax.numpy as jnp

a = jnp.array([1, 2, 3], dtype=jnp.bfloat16)
jnp.save('test.npy', a)
b = jnp.load('test.npy').view(jnp.bfloat16)
print(b)
# array([1, 2, 3], dtype=bfloat16)

We might think about doing this automatically within the jnp.load wrapper (currently jnp.load = np.load with no extra logic). Edit: made this change in #8499

[ivandustin]

Thank you. Is the pickle a different issue?

[ivandustin]

May I kindly also ask what is the recommended way or how do other projects save a model (a Python data structure that consists of JAX arrays) on the disk?

Currently, I have two options:

1. Pickle the entire Python data structure (e.g. a dict of dicts/lists) that consists of JAX arrays.
2. Manually jax.numpy.save each JAX array from a Python data structure to a single .npy file. Then manually jax.numpy.load each JAX array and rebuild the data structure.

The first one is easier, but I am not sure if there will be any problems when reading (pickle.loads) the file on another machine. (I had a problem with this in TensorFlow. If I pickled a tf.Variable that is located on an accelerator (e.g. TPU), if I load it on another machine without the same accelerator, an exception is thrown during pickle.loads. The solution is to copy the data first on the CPU before pickle.dumps.)

Thank you.

[grahamgower]

I've been using tolist() to convert types to their python equivalent and then serialise using whatever (I have small pytrees so will probably end up using json, as I don't care about the numpy/jax type when reading back in). E.g.

data = { ... } # some pytree
json.dumps(data, indent=2, default=lambda o: getattr(o, "tolist", lambda: o)())

But it looks like jax's tolist() doesn't behave as expected for bfloat16 --- it doesn't get converted to a python scalar.

import jax.numpy as jnp

a = jnp.zeros(5, dtype=jnp.bfloat16)
b = a.tolist()
print(type(a[0]), type(b[0]))


a = jnp.zeros(5, dtype=jnp.float16)
b = a.tolist()
print(type(a[0]), type(b[0]))

Output:

<class 'jaxlib.xla_extension.DeviceArray'> <class 'bfloat16'>
<class 'jaxlib.xla_extension.DeviceArray'> <class 'float'>

[jakevdp]

Regarding general use of pickle with JAX arrays: JAX delegates pickle to numpy. So any devicearray that you pickle will be saved (and re-loaded) as a standard numpy array, e.g.

import pickle
import jax.numpy as jnp
type(pickle.loads(pickle.dumps(jnp.arange(10))))
# numpy.ndarray

Regarding pickling of bfloat16: I think it's a deeper bug with how XLA defines the bfloat16 type. I opened #8505 to track the issue.