How does Jax return numpy arrays of bfloat16 dtype?

[botev]

I'm trying to understand, how does x.__arrray__() return a numpy array, but has a bfloat16 dtype, that generally do not exist in bare numpy?

[jakevdp]

The NumPy-compatible bfloat16 dtype is defined in the ml_dtypes package. You can use it directly like this:

import numpy as np
import ml_dtypes
np.arange(10, dtype=ml_dtypes.bfloat16)