How to test if a value is an jax dtype?

[samuela]

In PyTorch, one can use isinstance(thing, torch.dtype) to test whether thing is a torch.dtype. What's the analogous test in JAX?

Blindly using isinstance doesn't work:

In [11]: isinstance(jnp.int8, jnp.dtype)
Out[11]: False

[jakevdp]

You can use isinstance(thing, jnp.dtype) if you have a dtype; for example:

dt = jnp.dtype('int8')
isinstance(dt, jnp.dtype)  # True

Where it gets confusing is that jnp.int8 is not a dtype, it is a scalar constructor. JAX inherits this behavior from NumPy, where numpy.dtype('int8') is a dtype, while numpy.int8 is a scalar constructor which is not an instance of dtype.

If you are hoping to work explicitly with dtypes, you can always cast scalar constructors to dtypes using something like jnp.dtype(jnp.int8); if that doesn't help then perhaps you could share more about the problem you're trying to solve.

[samuela]

Ah, thank you for the informative answer @jakevdp! I was not aware of the dtype/constructor distinction