initial commit

geoopt/optim/radam.py

@@ -189,7 +189,7 @@ def perform_step(
         new_point, exp_avg_new = manifold.retr_transp(
             point, exp_avg, u=direction, t=-step_size
         )
-        point.set_(new_point)
+        point.set_(new_point.contiguous())
         exp_avg.set_(exp_avg_new)
 
     def stabilize_group(self, group):

geoopt/optim/rsgd.py

@@ -168,10 +168,10 @@ def perform_step(
                 point, momentum_buffer, u=grad, t=-lr
             )
             momentum_buffer.set_(new_momentum_buffer)
-            point.set_(new_point)
+            point.set_(new_point.contiguous())
         else:
             new_point = manifold.retr(point, grad, -lr)
-            point.set_(new_point)
+            point.set_(new_point.contiguous())
 
     def stabilize_group(self, group):
         with torch.no_grad():

geoopt/samplers/rhmc.py

@@ -40,7 +40,7 @@ def _step(self, p, r, epsilon):
 
         r.add_(epsilon * egrad2rgrad(p, p.grad))
         p_, r_ = retr_transp(p, r, u=r, t=epsilon)
-        p.set_(p_)
+        p.set_(p_.contiguous())
         r.set_(r_)
 
     def step(self, closure):

geoopt/samplers/rsgld.py

@@ -51,7 +51,7 @@ def step(self, closure):
                     n = torch.randn_like(p).mul_(math.sqrt(epsilon))
                     r = egrad2rgrad(p, 0.5 * epsilon * p.grad + n)
 
-                    p.set_(retr(p, r, 1.0))
+                    p.set_(retr(p, r, 1.0).contiguous())
                     p.grad.zero_()
 
         if not self.burnin:

geoopt/samplers/sgrhmc.py

@@ -76,7 +76,7 @@ def step(self, closure):
                         v = self.state[p]["v"]
 
                         p_, v_ = retr_transp(p, v, u=v, t=1.0)
-                        p.set_(p_)
+                        p.set_(p_.contiguous())
                         v.set_(v_)
 
                         n = egrad2rgrad(p, torch.randn_like(v))

geoopt/tensor.py

@@ -34,10 +34,10 @@ def proj_(self):
         Returns
         -------
         tensor
-            same instance
+            same instance (contiguous)
         """
         with torch.no_grad():
-            self.set_(self.manifold.projx(self.data))
+            self.set_(self.manifold.projx(self.data).contiguous())
         return self
 
     @insert_docs(Euclidean.retr.__doc__, r"\s+x : .+\n.+", "")

tests/test_adam.py

@@ -26,7 +26,7 @@ def closure():
         if (X - Xstar).norm() < 1e-5:
             break
         optim.step(closure)
-
+    assert X.is_contiguous()
     np.testing.assert_allclose(X.data, Xstar, atol=1e-5, rtol=1e-5)
     optim.load_state_dict(optim.state_dict())
     optim.step(closure)
@@ -49,5 +49,4 @@ def closure():
 
     for _ in range(2000):
         optim.step(closure)
-
     np.testing.assert_allclose(start.data, ideal, atol=1e-5, rtol=1e-5)

tests/test_rhmc.py

@@ -116,7 +116,7 @@ def forward(self):
 
     points = np.asarray(points)
     points = points[::20]
-
+    assert nd.x.is_contiguous()
     np.testing.assert_allclose(mu.numpy(), points.mean(axis=0), atol=1e-1)
     np.testing.assert_allclose(sigma.numpy(), points.std(axis=0), atol=1e-1)
 

tests/test_rsgd.py

@@ -34,7 +34,7 @@ def closure():
         if (X - Xstar).norm() < 1e-5:
             break
         optim.step(closure)
-
+    assert X.is_contiguous()
     np.testing.assert_allclose(X.data, Xstar, atol=1e-5)
     optim.load_state_dict(optim.state_dict())
     optim.step(closure)
@@ -56,5 +56,6 @@ def test_init_manifold():
     opt.zero_grad()
     opt.step()
     assert not np.allclose(p0.data, p0old.data)
+    assert p0.is_contiguous()
     np.testing.assert_allclose(p1.data, p1old.data)
     np.testing.assert_allclose(p0.data, stiefel.projx(p0old.data), atol=1e-4)