jax.experimental.disable_x64

jax.experimental.disable_x64#

jax.experimental.disable_x64()[source]#

Experimental context manager to temporarily disable 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:

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