rename invere root to root inverse

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/shampoo/shampoo.py

@@ -148,20 +148,18 @@ def _get_root_inverse(self, kronecker_factor: torch.Tensor) -> torch.Tensor:
         return (eigvecs * inv_root_scale) @ eigvecs.mT
 
     def precondition(self, x: torch.Tensor) -> torch.Tensor:
-        """Applies both inverse roots to a matrix in the parameter basis.
+        """Applies both root inverse to a matrix in the parameter basis.
 
         Args:
             x: Matrix in the parameter basis.
 
         Returns:
             The preconditioned matrix, in the parameter basis.
         """
-        inverse_root_pair = precond_base.TensorPair(
-            self._get_root_inverse(self.kronecker_factor_pair.L),
-            self._get_root_inverse(self.kronecker_factor_pair.R),
-        )
+        root_inv_L = self._get_root_inverse(self.kronecker_factor_pair.L)
+        root_inv_R = self._get_root_inverse(self.kronecker_factor_pair.R)
 
-        return inverse_root_pair.L @ x @ inverse_root_pair.R
+        return root_inv_L @ x @ root_inv_R
 
 
 class ShampooBase(optim.Optimizer, opt_mixin.WeightDecayMixin):