flax

1453d295 - Migrate from `jax.experimental.enable_x64` to `jax.enable_x64`.

Migrate from `jax.experimental.enable_x64` to `jax.enable_x64`. Also replace `jax.experimental.disable_x64()` with `jax.enable_x64(False)` The new public symbol is available in JAX v0.8.0, and the old experimental symbols are deprecated and will be removed in a future JAX version. PiperOrigin-RevId: 817376057