Graft using Adagrad. Essentially an implementation of Adagrad with momentum.

Partitions a tensor into smaller tensors for preconditioning. For example, if a variable has shape (4096, 512), we might split the 4096 into

Base class to perform grafting onto Shampoo. This class does no grafting.

Compute statistics/shape from gradients for preconditioning.

Graft using SGD+momentum. momentum maintains an exponentially weighted moving average of gradients.

Shampoo hyper parameters.

The Shampoo optimizer.