optax-shampoo (pypi): Distributed Shampoo (Second order Optimizer for Deep Learning) Optax Optimizer
Keywords: optax|distributed|shampoo|ai|machine-learning|research