What's the best way to get the floating point dtype?

[NeilGirdhar]

Right now, I'm doing:

@cache
def int_dtype() -> DTypeLike:
    return jnp.asarray(0).dtype


@cache
def float_dtype() -> DTypeLike:
    return jnp.empty(1).dtype


@cache
def complex_dtype() -> DTypeLike:
    return jnp.asarray(1j).dtype

[jakevdp]

Use dtypes.canonicalize_dtype:

from jax import dtypes
default_int = dtypes.canonicalize_dtype(int)
default_float = dtypes.canonicalize_dtype(float)
default_complex = dtypes.canonicalize_dtype(complex)

In general you should be careful about caching these values, because e.g. jax.experimental.enable_x64 can enable 64-bit values in a local context.

[NeilGirdhar]

Thanks for the lightning fast answer!

Right, the action of enable_x64 is what I'm trying to work around. I guess I can just call canonicalize in the appropriate places rather than having new functions.