flax

20446713 - 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). PiperOrigin-RevId: 900270194