search

google-deepmind / optax

Optax is a gradient processing and optimization library for JAX.
https://optax.readthedocs.io
Apache License 2.0
1.71k stars 194 forks source link

Support for Shampoo optimizer #1119

Closed kishorenc closed 3 weeks ago

kishorenc commented 1 month ago

It would be great if Optax can support Shampoo out of the box.

vroulet commented 3 weeks ago

Hello @kishorenc,

You can use https://github.com/google-research/google-research/tree/master/scalable_shampoo/optax.

vroulet commented 3 weeks ago

If the provided link does not match what you want please reopen the issue.