jax.experimental.disable_x64

jax.experimental.disable_x64()

Experimental context manager to temporarily disable X64 mode.

Warning

This context manager remains experimental because it is fundamentally broken and can result in unexpected behavior, particularly when used in conjunction with JAX transformations like jax.jit(), jax.vmap(), jax.grad(), and others. See jax-ml/jax#5982 for details.

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.