Merge pull request #2987 from IvyZX:push

PiperOrigin-RevId: 520124550

Merge pull request #2944 from IvyZX:push

PiperOrigin-RevId: 515489629

Merge pull request #2908 from IvyZX:push

PiperOrigin-RevId: 512756138

Merge pull request #2856 from levskaya:point_release4

PiperOrigin-RevId: 508228804

Merge pull request #2820 from IvyZX:push

PiperOrigin-RevId: 504950133

Merge pull request #2705 from IvyZX:push

PiperOrigin-RevId: 494237571

Merge pull request #2613 from IvyZX:push

PiperOrigin-RevId: 488547967

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

PiperOrigin-RevId: 478570635

Removes `flax.optim` & bumps version to 0.6.0

This fixes #2273

PiperOrigin-RevId: 467869152

Merge pull request #2313 from levskaya:idfix

PiperOrigin-RevId: 463175098