more efficient workaround

geoopt/optim/radam.py

@@ -4,6 +4,7 @@
 from .tracing import create_traced_update
 from ..tensor import ManifoldParameter, ManifoldTensor
 from ..manifolds import Euclidean
+from ..utils import copy_or_set
 
 
 class RiemannianAdam(OptimMixin, torch.optim.Adam):
@@ -190,7 +191,7 @@ def perform_step(
             point, exp_avg, u=direction, t=-step_size
         )
         # use copy only for user facing point
-        point.copy_(new_point)
+        copy_or_set(point, new_point)
         exp_avg.set_(exp_avg_new)
 
     def stabilize_group(self, group):
@@ -203,7 +204,7 @@ def stabilize_group(self, group):
                     continue
                 manifold = p.manifold
                 exp_avg = state["exp_avg"]
-                p.copy_(manifold.projx(p))
+                copy_or_set(p, manifold.projx(p))
                 exp_avg.set_(manifold.proju(p, exp_avg))
 
     def _sanitize_group(self, group):

geoopt/optim/rsgd.py

@@ -3,7 +3,7 @@
 from ..tensor import ManifoldParameter, ManifoldTensor
 from .mixin import OptimMixin
 from .tracing import create_traced_update
-
+from ..utils import copy_or_set
 
 __all__ = ["RiemannianSGD"]
 
@@ -169,10 +169,10 @@ def perform_step(
             )
             momentum_buffer.set_(new_momentum_buffer)
             # use copy only for user facing point
-            point.copy_(new_point)
+            copy_or_set(point, new_point)
         else:
             new_point = manifold.retr(point, grad, -lr)
-            point.copy_(new_point)
+            copy_or_set(point, new_point)
 
     def stabilize_group(self, group):
         with torch.no_grad():
@@ -181,7 +181,7 @@ def stabilize_group(self, group):
                     continue
                 manifold = p.manifold
                 momentum = group["momentum"]
-                p.copy_(manifold.projx(p))
+                copy_or_set(p, manifold.projx(p))
                 if momentum > 0:
                     param_state = self.state[p]
                     if not param_state:  # due to None grads

geoopt/samplers/rhmc.py

@@ -6,7 +6,7 @@
 from geoopt.tensor import ManifoldParameter, ManifoldTensor
 from geoopt.manifolds import Euclidean
 from geoopt.samplers.base import Sampler
-
+from ..utils import copy_or_set
 
 __all__ = ["RHMC"]
 
@@ -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.copy_(p_)
+        copy_or_set(p, p_)
         r.set_(r_)
 
     def step(self, closure):

geoopt/samplers/rsgld.py

@@ -5,7 +5,7 @@
 from geoopt.tensor import ManifoldParameter, ManifoldTensor
 from geoopt.manifolds import Euclidean
 from geoopt.samplers.base import Sampler
-
+from ..utils import copy_or_set
 
 __all__ = ["RSGLD"]
 
@@ -66,5 +66,4 @@ def stabilize(self):
                 for p in group["params"]:
                     if not isinstance(p, (ManifoldParameter, ManifoldTensor)):
                         continue
-
-                    p.copy_(p.manifold.projx(p))
+                    copy_or_set(p, p.manifold.projx(p))

geoopt/samplers/sgrhmc.py

@@ -5,7 +5,7 @@
 from geoopt.tensor import ManifoldParameter, ManifoldTensor
 from geoopt.manifolds import Euclidean
 from geoopt.samplers.base import Sampler
-
+from ..utils import copy_or_set
 
 __all__ = ["SGRHMC"]
 
@@ -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.copy_(p_)
+                        copy_or_set(p, p_)
                         v.set_(v_)
 
                         n = egrad2rgrad(p, torch.randn_like(v))
@@ -103,7 +103,6 @@ def stabilize(self):
 
                 manifold = p.manifold
                 v = self.state[p]["v"]
-
-                p.copy_(manifold.projx(p))
+                copy_or_set(p, manifold.projx(p))
                 # proj here is ok
                 v.set_(manifold.proju(p, v))

geoopt/utils.py

@@ -0,0 +1,23 @@
+__all__ = "copy_or_set"
+
+
+def copy_or_set(dest, source):
+    """
+    A workaround to respect strides of :code:`dest` when copying :code:`source`
+    (https://github.com/geoopt/geoopt/issues/70)
+
+    Parameters
+    ----------
+    dest : torch.Tensor
+        Destination tensor where to store new data
+    source : torch.Tensor
+        Source data to put in the new tensor
+    Returns
+    -------
+    dest
+        torch.Tensor
+    """
+    if dest.stride() != source.stride():
+        return dest.copy_(source)
+    else:
+        return dest.set_(source)