Merge pull request #4525 from google:v0.10.3

PiperOrigin-RevId: 725224019

Merge pull request #4387 from google:v0.10.2

PiperOrigin-RevId: 698095865

Merge pull request #4345 from google:v0.10.1

PiperOrigin-RevId: 691966843

Merge pull request #4292 from google:v0.10.0

PiperOrigin-RevId: 686679993

update Flax to version 0.9.0

PiperOrigin-RevId: 668047408

Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` u…

…seful.

Fixes #4008

PiperOrigin-RevId: 646679331

Merge pull request #3936 from google:nnx-stabilize

PiperOrigin-RevId: 636951986

Merge pull request #3876 from google:nnx-v0.1

PiperOrigin-RevId: 628712571

Fix tests after applying JAX key-reuse checker. See:

- https://jax.readthedocs.io/en/latest/jax.experimental.key_reuse.html#experimental-key-reuse-checking
- https://jax.readthedocs.io/en/latest/_autosummary/jax.random.clone.html#jax.random.clone

PiperOrigin-RevId: 615600399

Merge pull request #3682 from IvyZX:push

PiperOrigin-RevId: 605087622