Merge pull request google#4724 from google:0.10.6

PiperOrigin-RevId: 750649299

Merge pull request google#4656 from google:v0.10.5

PiperOrigin-RevId: 742270793

Merge pull request google#4579 from google:v0.10.4

PiperOrigin-RevId: 731490028

Merge pull request google#4525 from google:v0.10.3

PiperOrigin-RevId: 725224019

Merge pull request google#4387 from google:v0.10.2

PiperOrigin-RevId: 698095865

Merge pull request google#4345 from google:v0.10.1

PiperOrigin-RevId: 691966843

Merge pull request google#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 google#4008

PiperOrigin-RevId: 646679331

Merge pull request google#3936 from google:nnx-stabilize

PiperOrigin-RevId: 636951986