Files copied from google-research/scalable_shampoo/optax.
Imports have been modified to be relative.
This will eventually be replaced with optax-shampoo package.
optax-shampoo
