Need more numerical range or decimal precision than float32 provides? float64 is a 64-bit IEEE 754 floating-point type with roughly 15–17 significant decimal digits and a range near 2.23e-308 to 1.80e308. JAX defaults to float32, so requesting float64 can silently produce float32 unless 64-bit mode is enabled. section 3.3.2
Enable float64 computations
Enable 64-bit mode at startup, then create the array with dtype=jnp.float64. Check x.dtype to confirm the result. The example notes that this may not work on TPU backends, so CPU or GPU may be required. section 3.3.2
Listing 3.24 Forcing float64 computations
# this only works on startup!
from jax.config import config
config.update("jax_enable_x64", True) #1
import jax.numpy as jnp
# this may not work on TPU backend. Try using CPU or GPU.
x = jnp.array(range(10), dtype=jnp.float64) #2
x.dtype
>>> dtype('float64') #3
Shows the startup configuration, array construction, and dtype check needed to force float64 computations. section 3.3.2
