Run AdaBound on the real and imaginary parts

tests/test_adabound_complex.py

@@ -0,0 +1,68 @@
+"""Complex AdaBound is the real algorithm on both components."""
+import torch
+import torch_optimizer as optim
+
+BASE = torch.tensor(
+    [
+        [0.1, -0.2, 0.3],
+        [0.0, 1.5, -0.4],
+        [2.0, -1.0, 0.5],
+        [0.7, 0.7, -0.7],
+    ]
+)
+GRAD = torch.tensor(
+    [
+        [0.2, 0.0, -0.1],
+        [1.0, -0.5, 0.3],
+        [0.0, 0.4, -0.2],
+        [0.1, 0.1, 0.1],
+    ]
+)
+
+
+def test_real_path_is_unchanged():
+    param = torch.nn.Parameter(BASE.clone())
+    opt = optim.AdaBound([param], lr=1e-3, weight_decay=1e-2)
+    param.grad = GRAD.clone()
+    opt.step()
+    param.grad = (-0.5 * GRAD).clone()
+    opt.step()
+    expected = torch.tensor(
+        [
+            [0.09872793406248093, -0.19800041615962982, 0.30123090744018555],
+            [-0.0012663287343457341, 1.501231074333191, -0.4012507498264313],
+            [1.998000144958496, -1.0012370347976685, 0.5012369155883789],
+            [0.6986559629440308, 0.6986559629440308, -0.7011821269989014],
+        ]
+    )
+    assert torch.equal(param.detach(), expected)
+
+
+def _complex_matches(amsbound):
+    base = torch.tensor([1.0 + 0.5j, -0.25 + 0.0j, 0.3 - 0.8j])
+    grad = torch.tensor([0.4 - 0.1j, 0.2 + 0.2j, -0.5 + 0.3j])
+    param = torch.nn.Parameter(base.clone())
+    opt = optim.AdaBound(
+        [param], lr=1e-2, betas=(0.9, 0.999), final_lr=0.1, amsbound=amsbound
+    )
+    param.grad = grad.clone()
+    opt.step()
+    real_param = torch.nn.Parameter(torch.view_as_real(base).clone())
+    real_opt = optim.AdaBound(
+        [real_param],
+        lr=1e-2,
+        betas=(0.9, 0.999),
+        final_lr=0.1,
+        amsbound=amsbound,
+    )
+    real_param.grad = torch.view_as_real(grad).clone()
+    real_opt.step()
+    assert torch.equal(torch.view_as_real(param.detach()), real_param.detach())
+
+
+def test_complex_matches_real_algorithm_on_both_parts():
+    _complex_matches(False)
+
+
+def test_amsbound_complex_matches_real_algorithm():
+    _complex_matches(True)

torch_optimizer/adabound.py

@@ -112,6 +112,13 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                         "please consider SparseAdam instead"
                     )
                     raise RuntimeError(msg)
+                # The bound clamp is a real op. A complex parameter is the
+                # same algorithm on its real and imaginary parts, which is
+                # how Adam treats complex weights. Real parameters stay on
+                # the path below.
+                complex_param = torch.is_complex(p)
+                if complex_param:
+                    grad = torch.view_as_real(grad)
                 amsbound = group["amsbound"]
 
                 state = self.state[p]
@@ -137,12 +144,20 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                 exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                 if amsbound:
                     max_exp_avg_sq = state["max_exp_avg_sq"]
+                if complex_param:
+                    exp_avg = torch.view_as_real(exp_avg)
+                    exp_avg_sq = torch.view_as_real(exp_avg_sq)
+                    if amsbound:
+                        max_exp_avg_sq = torch.view_as_real(max_exp_avg_sq)
                 beta1, beta2 = group["betas"]
 
                 state["step"] += 1
 
                 if group["weight_decay"] != 0:
-                    grad = grad.add(p.data, alpha=group["weight_decay"])
+                    decay_base = (
+                        torch.view_as_real(p.data) if complex_param else p.data
+                    )
+                    grad = grad.add(decay_base, alpha=group["weight_decay"])
 
                 # Decay the first and second moment running average coefficient
                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
@@ -179,5 +194,8 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                     exp_avg
                 )
 
-                p.data.add_(-step_size)
+                if complex_param:
+                    torch.view_as_real(p.data).add_(-step_size)
+                else:
+                    p.data.add_(-step_size)
         return loss