jax.experimental.enable_x64

jax.experimental.enable_x64(new_val=True)

Experimental context manager to temporarily enable X64 mode.

Warning

This context manager is deprecated as of JAX v0.8.0, and will be removed in JAX v0.9.0. Use jax.enable_x64() instead.

Usage:

>>> import jax
>>> x = np.arange(5, dtype='float64')
>>> with _enable_x64(True):
...   print(jnp.asarray(x).dtype)
...
float64

See also

jax.experimental.disable_x64

temporarily disable X64 mode.

Parameters:

new_val (bool)