From ac6a5002b50f7f6e6e33771d0c9542b25087b9f9 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Wed, 14 Oct 2026 11:04:52 -0700 Subject: [PATCH 1/2] change init value Signed-off-by: Hao Wu --- emerging_optimizers/shampoo/shampoo.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/emerging_optimizers/shampoo/shampoo.py b/emerging_optimizers/shampoo/shampoo.py index a161f922..7aa06ec5 100644 --- a/emerging_optimizers/shampoo/shampoo.py +++ b/emerging_optimizers/shampoo/shampoo.py @@ -195,8 +195,10 @@ def init_state( { "Q_L": torch.eye(m, device=device), "Q_R": torch.eye(n, device=device), - "eigvals_L": torch.zeros(m, device=device), - "eigvals_R": torch.zeros(n, device=device), + # Initialize eigen values one to skip KL correction for the first step. + # They'll be override right after + "eigvals_L": torch.ones(m, device=device), + "eigvals_R": torch.ones(n, device=device), } ) return state From 070a043f1d623f4aa1b75fffb9b4d4189f376d53 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Wed, 14 Oct 2026 11:16:37 -0700 Subject: [PATCH 2/2] grammar fix Signed-off-by: Hao Wu --- emerging_optimizers/shampoo/shampoo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/emerging_optimizers/shampoo/shampoo.py b/emerging_optimizers/shampoo/shampoo.py index 7aa06ec5..850f7f8c 100644 --- a/emerging_optimizers/shampoo/shampoo.py +++ b/emerging_optimizers/shampoo/shampoo.py @@ -195,8 +195,8 @@ def init_state( { "Q_L": torch.eye(m, device=device), "Q_R": torch.eye(n, device=device), - # Initialize eigen values one to skip KL correction for the first step. - # They'll be override right after + # Initialize eigenvalues to one to skip KL correction for the first step. + # They will be overwritten immediately afterward. "eigvals_L": torch.ones(m, device=device), "eigvals_R": torch.ones(n, device=device), }