Our implementation of Shampoo optimizer based on https://arxiv.org/pdf/1802.09568.pdf. It consists of different notebooks, which we used on our own computers or Google Colab.
Use scripts matrix_square_root_power.py and shampoo_optimizer.py both for using our code. Method apply_gradients does one iteration of Shampoo optimization process.