move to linalg

geoopt/__init__.py

@@ -2,7 +2,7 @@
 from . import optim
 from . import tensor
 from . import samplers
-from . import util
+from . import linalg
 
 from .tensor import ManifoldParameter, ManifoldTensor
 from .manifolds import Stiefel, Euclidean, Sphere

geoopt/linalg/__init__.py

@@ -0,0 +1 @@
+from .batch_linalg import svd, qr, sym, extract_diag, matrix_rank

geoopt/manifolds/sphere.py

@@ -1,7 +1,7 @@
 import torch
 
 from .base import Manifold
-import geoopt.util.linalg
+import geoopt.linalg.batch_linalg
 
 __all__ = [
     "Sphere",
@@ -95,7 +95,7 @@ class SphereSubspaceIntersection(Sphere):
 
     def __init__(self, span):
         self._configure_manifold(span)
-        if (geoopt.util.linalg.matrix_rank(self._projector) == 1).any():
+        if (geoopt.linalg.batch_linalg.matrix_rank(self._projector) == 1).any():
             raise ValueError(
                 "Manifold only consists of isolated points when "
                 "subspace is 1-dimensional."
@@ -119,7 +119,7 @@ def _check_shape(self, x, name):
         return ok, reason
 
     def _configure_manifold(self, span):
-        Q, _ = geoopt.util.linalg.qr(span)
+        Q, _ = geoopt.linalg.batch_linalg.qr(span)
         self._projector = Q @ Q.transpose(-1, -2)
 
     def _project_on_subspace(self, x):
@@ -150,7 +150,7 @@ class SphereSubspaceComplementIntersection(SphereSubspaceIntersection):
     """
 
     def _configure_manifold(self, span):
-        Q, _ = geoopt.util.linalg.qr(span)
+        Q, _ = geoopt.linalg.batch_linalg.qr(span)
         P = -Q @ Q.transpose(-1, -2)
         P[..., torch.arange(P.shape[-2]), torch.arange(P.shape[-2])] += 1
         self._projector = P

geoopt/manifolds/stiefel.py

@@ -1,6 +1,6 @@
 import torch
 
-from geoopt import util
+from geoopt import linalg
 from .base import Manifold
 
 
@@ -74,7 +74,7 @@ def _amat(self, x, u):
         return u @ x.transpose(-1, -2) - x @ u.transpose(-1, -2)
 
     def _projx(self, x):
-        U, d, V = util.linalg.svd(x)
+        U, d, V = linalg.batch_linalg.svd(x)
         return torch.einsum("...ik,...k,...jk->...ij", [U, torch.ones_like(d), V])
 
 
@@ -152,7 +152,7 @@ class EuclideanStiefel(Stiefel):
     reversible = False
 
     def _proju(self, x, u):
-        return u - x @ util.linalg.sym(x.transpose(-1, -2) @ u)
+        return u - x @ linalg.batch_linalg.sym(x.transpose(-1, -2) @ u)
 
     def _transp_one(self, x, u, t, v, y=None):
         if y is None:
@@ -173,7 +173,7 @@ def _inner(self, x, u, v):
         return (u * v).sum([-1, -2])
 
     def _retr(self, x, u, t):
-        q, r = util.linalg.qr(x + u * t)
-        unflip = torch.sign(torch.sign(util.linalg.extract_diag(r)) + 0.5)
+        q, r = linalg.batch_linalg.qr(x + u * t)
+        unflip = torch.sign(torch.sign(linalg.batch_linalg.extract_diag(r)) + 0.5)
         q *= unflip[..., None, :]
         return q

tests/test_utils.py

@@ -0,0 +1,17 @@
+import pytest
+import torch
+import numpy as np
+import geoopt
+
+
+@pytest.fixture
+def A():
+    torch.manual_seed(42)
+    n = 10
+    a = torch.randn(n, 3, 3)
+    a[:, 2, :] = 0
+    return a
+
+
+def test_svd(A):
+    u, d, v = geoopt.linalg.batch_linalg.svd(A)