Skip redundant full-rank decompositions (#52)

Author: lezcano

geotorch/lowrank.py

@@ -149,6 +149,10 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
         with torch.no_grad():
             X = self[0].base.new_empty(self.tensorial_size + (self.n, self.k))
             init_(X)
+            if self.rank == self.k and not factorized:
+                if self.transposed:
+                    X = X.transpose(-2, -1)
+                return X
             U, S, Vt = torch.linalg.svd(X, full_matrices=False)
             U, S, Vt = U[..., : self.rank], S[..., : self.rank], Vt[..., : self.rank, :]
             if factorized:

geotorch/symmetric.py

@@ -218,6 +218,8 @@ def sample(self, init_=torch.nn.init.xavier_normal_, factorized=False):
             X = self[0].base.new_empty(self.tensorial_size + (self.n, self.n))
             init_(X)
             X = X @ X.transpose(-2, -1)
+            if self.rank == self.n and not factorized:
+                return X
             L, Q = torch.linalg.eigh(X)
             L = L[..., -self.rank :]
             Q = Q[..., -self.rank :]

test/test_lowrank.py

@@ -1,4 +1,4 @@
-from unittest import TestCase
+from unittest import TestCase, mock
 
 import torch
 
@@ -35,3 +35,24 @@ def test_inverse_softplus_is_stable(self):
         inverse = inv_softplus_epsilon(x)
         self.assertTrue(torch.isfinite(inverse).all())
         self.assertTrue(torch.allclose(softplus_epsilon(inverse), x))
+
+    def test_full_rank_sample_skips_svd(self):
+        for size in [(3, 3), (5, 3), (3, 5), (2, 5, 3)]:
+            manifold = LowRank(size=size, rank=min(size[-2:])).double()
+
+            def init_(X):
+                return X.copy_(
+                    torch.arange(X.numel(), dtype=X.dtype, device=X.device).view_as(X)
+                )
+
+            canonical_size = size[:-2] + (max(size[-2:]), min(size[-2:]))
+            expected = torch.empty(canonical_size, dtype=torch.float64)
+            expected = torch.arange(expected.numel(), dtype=expected.dtype).view_as(
+                expected
+            )
+            if size[-2] < size[-1]:
+                expected = expected.mT
+
+            with mock.patch("torch.linalg.svd", side_effect=AssertionError):
+                sample = manifold.sample(init_)
+            self.assertTrue(torch.equal(sample, expected))

test/test_symmetric.py

@@ -1,12 +1,13 @@
 # Tests for the Sphere
-from unittest import TestCase
+from unittest import TestCase, mock
 import itertools
 
 import torch
 import torch.nn as nn
 from torch.nn.utils import parametrize
 
 from geotorch.symmetric import Symmetric, SymF
+from geotorch.pssd import PSSD
 
 
 class TestSymmetric(TestCase):
@@ -64,3 +65,17 @@ def test_construction(self):
 
     def test_repr(self):
         print(Symmetric())
+
+    def test_full_rank_sample_skips_eigendecomposition(self):
+        manifold = PSSD(size=(2, 4, 4)).double()
+
+        def init_(X):
+            return X.copy_(
+                torch.arange(X.numel(), dtype=X.dtype, device=X.device).view_as(X)
+            )
+
+        X = torch.arange(32, dtype=torch.float64).view(2, 4, 4)
+        expected = X @ X.mT
+        with mock.patch("torch.linalg.eigh", side_effect=AssertionError):
+            sample = manifold.sample(init_)
+        self.assertTrue(torch.equal(sample, expected))