Update some test shmoo

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

Author: skyw

emerging_optimizers/soap/soap.py

@@ -89,7 +89,7 @@ def __init__(
         self,
         params: ParamsT,
         lr: float,
-        betas: Tuple[float, float] = (0.95, 0.9),
+        betas: Tuple[float, float] = (0.9, 0.95),
         shampoo_beta: float = 0.95,
         eps: float = 1e-8,
         weight_decay: float = 0.01,

tests/soap_mnist_test.py

@@ -47,14 +47,10 @@ def forward(self, x):
     "eps": 1e-8,
     "precondition_1d": True,  # Enable preconditioning for bias vectors
     "precondition_frequency": 1,  # Update preconditioner every step for testing
-    "trace_normalization": True,
-    "shampoo_beta": 0.9,  # Slightly more aggressive moving average
     "fp32_matmul_prec": "high",
     "qr_fp32_matmul_prec": "high",
     "use_adaptive_criteria": False,
     "power_iter_steps": 1,
-    "use_nesterov": True,
-    "skip_preconditioning_steps": 0,
 }
 
 
@@ -106,15 +102,12 @@ def main() -> None:
     # Initialize optimizers
     optimizer_soap = SOAP(
         model_soap.parameters(),
-        lr=9.0 * config["lr"],
+        lr=2.1 * config["lr"],
         weight_decay=config["weight_decay"],
         betas=(config["adam_beta1"], config["adam_beta2"]),
         eps=config["eps"],
         precondition_frequency=config["precondition_frequency"],
-        trace_normalization=config["trace_normalization"],
-        shampoo_beta=config["shampoo_beta"],
         precondition_1d=config["precondition_1d"],
-        use_nesterov=config["use_nesterov"],
         fp32_matmul_prec=config["fp32_matmul_prec"],
         qr_fp32_matmul_prec=config["qr_fp32_matmul_prec"],
         use_adaptive_criteria=config["use_adaptive_criteria"],