jax.experimental.disable_x64

jax.experimental.disable_x64()

Experimental context manager to temporarily disable X64 mode.

Usage:

>>> x = np.arange(5, dtype='float64')
>>> with disable_x64():
...   print(jnp.asarray(x).dtype)
...
float32

See also

jax.experimental.enable_x64

temporarily enable X64 mode.