flax

f6109649 - Avoid usage of deprecated `jax.core` APIs.

Avoid usage of deprecated `jax.core` APIs. These APIs are deprecated as of JAX v0.10.0, replaced by equivalents in `jax.extend.core` (see https://docs.jax.dev/en/latest/jax.extend.html for details). Bumps jax to `>=0.10.0`. PiperOrigin-RevId: 901539317