Saving JAX Tensors of dtype bfloat16 with safetensors

[gokul-uf]

Hi folks,
The flax binding (flax.py) depends on np.asarray to convert JAX tensors to NumPy here.

NumPy does not support bfloat16 (https://numpy.org/doc/stable/user/basics.types.html) but somehow the np.asarray call passes and give us the impression that we can go from JAX -> Numpy at bfloat16.

I haven't run any tests yet of numerical impacts of this but wanted to raise this gotcha and see if folks think this could be a silent bug.

Thanks!