Merge 56a7c61 into d15832a

geoopt/samplers/rhmc.py

@@ -148,3 +148,14 @@ def step(self, closure):
                 self.log_probs.append(old_logp)
             else:
                 self.log_probs.append(new_logp)
+
+    @torch.no_grad()
+    def stabilize_group(self, group):
+        for p in group["params"]:
+            if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
+                continue
+            copy_or_set_(p, p.manifold.projx(p))
+            state = self.state[p]
+            if not state:  # due to None grads
+                continue
+            copy_or_set_(state["old_p"], p.manifold.projx(state["old_p"]))

geoopt/samplers/rsgld.py

@@ -57,12 +57,9 @@ def step(self, closure):
             self.steps += 1
             self.log_probs.append(logp.item())
 
-    def stabilize(self):
-        """Stabilize parameters if they are off-manifold due to numerical reasons
-        """
-        with torch.no_grad():
-            for group in self.param_groups:
-                for p in group["params"]:
-                    if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
-                        continue
-                    copy_or_set_(p, p.manifold.projx(p))
+    @torch.no_grad()
+    def stabilize_group(self, group):
+        for p in group["params"]:
+            if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
+                continue
+            copy_or_set_(p, p.manifold.projx(p))

geoopt/samplers/sgrhmc.py

@@ -91,17 +91,16 @@ def step(self, closure):
             self.steps += 1
             self.log_probs.append(logp)
 
-    def stabilize(self):
-        """Stabilize parameters if they are off-manifold due to numerical reasons
-        """
-
-        for group in self.param_groups:
-            for p in group["params"]:
-                if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
-                    continue
-
-                manifold = p.manifold
-                v = self.state[p]["v"]
-                copy_or_set_(p, manifold.projx(p))
-                # proj here is ok
-                v.set_(manifold.proju(p, v))
+    @torch.no_grad()
+    def stabilize_group(self, group):
+        for p in group["params"]:
+            if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
+                continue
+
+            manifold = p.manifold
+            copy_or_set_(p, manifold.projx(p))
+            # proj here is ok
+            state = self.state[p]
+            if not state:
+                continue
+            state["v"].set_(manifold.proju(p, state["v"]))