Fix LookSAM gv state key mismatch and float alpha handling

src/infty/optim/geometry_reshaping/looksam.py

@@ -38,9 +38,10 @@ def sam_perturb_weights(self):
     def compute_gv(self):
         for group in self.param_groups:
             for p in group['params']:
-                if p.grad is None or 'old_grad' not in self.state[p]:
+                # sam_perturb_weights stores the first-order grad under 'old_g'
+                if p.grad is None or 'old_g' not in self.state[p]:
                     continue
-                old_grad = self.state[p]['old_grad']
+                old_grad = self.state[p]['old_g']
                 new_grad = p.grad
                 g_grad_norm = old_grad / (old_grad.norm(p=2) + 1e-8)
                 g_new_grad_norm = new_grad / (new_grad.norm(p=2) + 1e-8)
@@ -55,7 +56,8 @@ def apply_gv(self):
                 if p.grad is None or 'gv' not in self.state[p]:
                     continue
                 gv = self.state[p]['gv']
-                p.grad.add_(self.alpha.to(p) * (p.grad.norm(p=2) / (gv.norm(p=2) + 1e-8)) * gv)
+                alpha = self.alpha.to(p) if torch.is_tensor(self.alpha) else self.alpha
+                p.grad.add_(alpha * (p.grad.norm(p=2) / (gv.norm(p=2) + 1e-8)) * gv)
 
     def step(self, closure=None, delay=False):
         if closure:

tests/optim/test_looksam_gv.py

@@ -0,0 +1,66 @@
+import torch
+
+
+def test_looksam_stores_gv_after_periodic_sam_step():
+    """LookSAM must persist a vertical component (`gv`) after each k-step SAM refresh.
+
+    Regression: compute_gv looked up state['old_grad'] while sam_perturb_weights
+    stores state['old_g'], so gv was never written and non-k steps silently
+    skipped apply_gv (LookSAM degraded to plain SGD between refreshes).
+    """
+    from infty.optim.geometry_reshaping.looksam import LookSAM
+
+    torch.manual_seed(0)
+    model = torch.nn.Linear(4, 2)
+    base_optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    optimizer = LookSAM(
+        model.parameters(),
+        base_optimizer=base_optimizer,
+        model=model,
+        args={"rho": 0.05, "k": 2, "alpha": 0.7},
+    )
+
+    x = torch.randn(6, 4)
+    y = torch.tensor([0, 1, 0, 1, 0, 1])
+    criterion = torch.nn.CrossEntropyLoss()
+
+    def closure():
+        logits = model(x)
+        loss = criterion(logits, y)
+        return logits, [loss]
+
+    optimizer.set_closure(closure)
+
+    optimizer.step()  # step 1: non-k, no gv yet
+    optimizer.step()  # step 2: k-multiple -> SAM + compute_gv
+
+    assert any(
+        "gv" in optimizer.state.get(p, {}) for p in model.parameters()
+    ), "LookSAM.compute_gv did not store 'gv' after a periodic SAM step"
+
+
+def test_looksam_apply_gv_accepts_float_alpha():
+    """apply_gv must accept Python-float alpha (default 0.7), not call alpha.to(p)."""
+    from infty.optim.geometry_reshaping.looksam import LookSAM
+
+    torch.manual_seed(1)
+    model = torch.nn.Linear(4, 2)
+    base_optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    optimizer = LookSAM(
+        model.parameters(),
+        base_optimizer=base_optimizer,
+        model=model,
+        args={"rho": 0.05, "k": 2, "alpha": 0.7},
+    )
+    assert isinstance(optimizer.alpha, float)
+
+    for p in model.parameters():
+        p.grad = torch.ones_like(p)
+        optimizer.state[p]["gv"] = torch.randn_like(p)
+
+    # Must not raise AttributeError: 'float' object has no attribute 'to'
+    optimizer.apply_gv()
+
+    for p in model.parameters():
+        assert p.grad is not None
+        assert torch.isfinite(p.grad).all()