refactor: AdaGradGraft

Author: kozistr

pytorch_optimizer/__init__.py

@@ -41,6 +41,21 @@
 from pytorch_optimizer.optimizer.sam import SAM
 from pytorch_optimizer.optimizer.sgdp import SGDP
 from pytorch_optimizer.optimizer.shampoo import Shampoo
+from pytorch_optimizer.optimizer.shampoo_utils import (
+    AdaGradGraft,
+    BlockPartitioner,
+    Graft,
+    LayerWiseGrafting,
+    PreConditioner,
+    PreConditionerType,
+    RMSPropGraft,
+    SGDGraft,
+    SQRTNGraft,
+    compute_power,
+    matrix_power,
+    merge_small_dims,
+    power_iter,
+)
 from pytorch_optimizer.optimizer.utils import (
     clip_grad_norm,
     disable_running_stats,

pytorch_optimizer/optimizer/shampoo.py

@@ -5,7 +5,7 @@
 from pytorch_optimizer.base.optimizer import BaseOptimizer
 from pytorch_optimizer.base.types import BETAS, CLOSURE, DEFAULTS, LOSS, PARAMETERS
 from pytorch_optimizer.optimizer.shampoo_utils import (
-    AdagradGraft,
+    AdaGradGraft,
     Graft,
     LayerWiseGrafting,
     PreConditioner,
@@ -130,7 +130,7 @@ def reset(self):
                     self.pre_conditioner_type,
                 )
                 if self.graft_type == LayerWiseGrafting.ADAGRAD:
-                    state['graft'] = AdagradGraft(p, self.diagonal_eps)
+                    state['graft'] = AdaGradGraft(p, self.diagonal_eps)
                 elif self.graft_type == LayerWiseGrafting.RMSPROP:
                     state['graft'] = RMSPropGraft(p, self.diagonal_eps)
                 elif self.graft_type == LayerWiseGrafting.SGD:
@@ -172,7 +172,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                         self.pre_conditioner_type,
                     )
                     if self.graft_type == LayerWiseGrafting.ADAGRAD:
-                        state['graft'] = AdagradGraft(p, self.diagonal_eps)
+                        state['graft'] = AdaGradGraft(p, self.diagonal_eps)
                     elif self.graft_type == LayerWiseGrafting.RMSPROP:
                         state['graft'] = RMSPropGraft(p, self.diagonal_eps)
                     elif self.graft_type == LayerWiseGrafting.SGD:

pytorch_optimizer/optimizer/shampoo_utils.py

@@ -66,7 +66,7 @@ def precondition_gradient(self, grad: torch.Tensor) -> torch.Tensor:
         return torch.ones_like(grad) * torch.sign(grad)
 
 
-class AdagradGraft(SGDGraft):
+class AdaGradGraft(SGDGraft):
     r"""Graft using Adagrad. Essentially an implementation of Adagrad with momentum.
 
     :param var: torch.Tensor. variable.