Is it safe to use JAX dtypes for type casting in numpy?

[uduse]

Curious about the best practices on doing type casting in either numpy arrays or jax numpy arrays. How compatible it is for numpy array to handle jax dtypes? i.e., can I use JAX types everywhere for both numpy array handlings and jax numpy array handlings?

e.g., would np.array(values, dtype=jnp.float16) ever cause problem? What happens when we do something like np.array(values, dtype=jnp.bfloat16)? The code runs but I'm not sure if there's dire consequences, since CPUs shouldn't have bfloat16?

[jakevdp]

Short answer, yes!

Long answer: jnp.float16 is not a dtype, it's a JAX scalar array constructor. This is a mirror of the design decision of NumPy, where np.float16 is not a dtype, but rather a scalar object constructor. For convenience, NumPy lets you write things like dtype=np.float16 as a shorthand for dtype=np.dtype('float16'), and we wanted dtype=jnp.float16 to work similarly.

NumPy scalar types like np.float16 are recognized as valid dtypes due to specific code paths in the np.dtype constructor designed to support this shortcut.

JAX scalar types like jnp.float16 are recognized as valid dtypes because np.dtype looks for a dtype attribute on any unknown object passed to it, and JAX scalar constructors like define a suitable dtype attribute specifically to support this use-case.

since CPUs shouldn't have bfloat16?

CPUs do have bfloat16: the bfloat16 dtype is defined in the ml_dtypes package, and this dtype can be used without ever importing JAX:

In [1]: import ml_dtypes

In [2]: import numpy as np

In [3]: np.arange(10, dtype=ml_dtypes.bfloat16)
Out[3]: array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=bfloat16)

This is the same bfloat16 dtype that is exposed in JAX. That said, not all NumPy or SciPy operations are defined for this bfloat16 dtype, so if you want to do any nontrivial operations on these arrays, you should use JAX.