cannot pickle bfloat16 array

[guolinke]

Please:

• Check for duplicate issues.
• Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

>>> import jax.numpy as jnp
>>> a = [1, 3, 3, 5]
>>> t = jnp.asarray(a, dtype=jnp.bfloat16)
>>> import pickle
>>> pp = pickle.dump(t, file=open('test.pkl','wb'))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
_pickle.PicklingError: Can't pickle <class 'bfloat16'>: attribute lookup bfloat16 on builtins failed

>>> import jax
>>> jax.__version__
'0.2.25'
>>> import jaxlib
>>> jaxlib.__version__
'0.1.73'

• If applicable, include full error messages/tracebacks.

[jakevdp]

Hi - thanks for the report! This issue is tracked in #8505, so I'm going to mark this as a duplicate.