Merge pull request google#2856 from levskaya:point_release4

PiperOrigin-RevId: 508228804

Merge pull request google#2820 from IvyZX:push

PiperOrigin-RevId: 504950133

Merge pull request google#2705 from IvyZX:push

PiperOrigin-RevId: 494237571

Merge pull request google#2613 from IvyZX:push

PiperOrigin-RevId: 488547967

Merge pull request google#2435 from jheek:boxed-metadata-flip

PiperOrigin-RevId: 478570635

Removes `flax.optim` & bumps version to 0.6.0

This fixes google#2273

PiperOrigin-RevId: 467869152

Merge pull request google#2313 from levskaya:idfix

PiperOrigin-RevId: 463175098

Merge pull request google#2203 from cgarciae:v0.5.2

PiperOrigin-RevId: 456237317

Merge pull request google#2182 from yechengxi:patch-1

PiperOrigin-RevId: 454138229

Merge pull request google#2134 from jheek:version-0.5.0

PiperOrigin-RevId: 450415939