[Feature] Implement more Shampoo features

[kozistr]

Problem (Why?)

the original JAX implementation has more features.

Solution (What/How?)

• implement moving_average_for_momentum
• implement decoupled_weight_decay
• implement decoupled_learning_rate
• implement more grafting types
  • RMSProp
  • SQRT_N
• implement type of preconditioner
  • ALL
  • INPUT (one-sided Shampoo)

Other changes (bug fixes, small refactors)

nope

Notes

• still there's a limitation on distributed training.
• bump version to v2.3.1

[codecov-commenter]

Codecov Report

Merging #99 (5da4074) into main (40f34df) will increase coverage by 0.00%.
The diff coverage is 100.00%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

@@           Coverage Diff           @@
##             main      #99   +/-   ##
=======================================
  Coverage   99.64%   99.65%           
=======================================
  Files          37       37           
  Lines        2568     2618   +50     
=======================================
+ Hits         2559     2609   +50     
  Misses          9        9           

Impacted Files	Coverage Δ
pytorch_optimizer/optimizer/shampoo.py	100.00% <100.00%> (ø)
pytorch_optimizer/optimizer/shampoo_utils.py	100.00% <100.00%> (ø)
pytorch_optimizer/optimizer/utils.py	97.75% <100.00%> (+0.02%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.