Merge pull request google#2987 from IvyZX:push

PiperOrigin-RevId: 520124550

Merge pull request google#2944 from IvyZX:push

PiperOrigin-RevId: 515489629

Merge pull request google#2908 from IvyZX:push

PiperOrigin-RevId: 512756138

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