Support SOAP optimizer

[skyw]

Is your feature request related to a problem? Please describe.

SOAP (https://arxiv.org/abs/2409.11321) is a preconditioned optimizer that combines the benefits of Shampoo’s non-diagonal preconditioning with Adam’s adaptive learning rates

Describe the solution you'd like

A clear and concise description of what you want to happen.
Provide a code snippet on how new APIs/changes would be used by others.

Describe alternatives you've considered

A clear and concise description of any alternative solutions or features you've considered.

Additional context

Add any other context or screenshots about the feature request here