Reduce manifold sampling temporaries

Author: lezcano

geotorch/fixedrank.py

@@ -113,8 +113,8 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6, factorized=False)
         """
         U, S, V = super().sample(factorized=True, init_=init_)
         with torch.no_grad():
-            # S >= 0, as given by torch.linalg.eigvalsh()
-            S[S < eps] = eps
+            # S >= 0, as given by torch.linalg.svd()
+            S.clamp_min_(eps)
         if factorized:
             return U, S, V
         else:

geotorch/lowrank.py

@@ -107,7 +107,7 @@ def in_manifold_singular_values(self, S, eps=1e-5):
             return True
         # We compute the \infty-norm of the remaining dimension
         D = S[..., self.rank :]
-        infty_norm_err = D.abs().max(dim=-1).values
+        infty_norm_err = D.abs().amax(dim=-1)
         return (infty_norm_err < eps).all()
 
     def in_manifold(self, X, eps=1e-5):

geotorch/pssdfixedrank.py

@@ -95,5 +95,5 @@ def sample(self, init_=torch.nn.init.xavier_normal_, eps=5e-6):
         L, Q = super().sample(factorized=True, init_=init_)
         with torch.no_grad():
             # L >= 0, as given by torch.linalg.eigvalsh()
-            L[L < eps] = eps
+            L.clamp_min_(eps)
         return (Q * L.unsqueeze(-2)) @ Q.transpose(-2, -1)

geotorch/sphere.py

@@ -16,14 +16,13 @@ def uniform_init_sphere_(x, r=1.0):
     """
     with torch.no_grad():
         x.normal_()
-        x.copy_(r * project(x))
+        x.copy_(project(x)).mul_(r)
     return x
 
 
 def _in_sphere(x, r, eps):
     norm = torch.linalg.vector_norm(x, dim=-1)
-    rs = torch.full_like(norm, r)
-    return (torch.linalg.vector_norm(norm - rs, ord=float("inf")) < eps).all()
+    return (norm - r).abs().amax() < eps
 
 
 class SphereEmbedded(nn.Module):

geotorch/symmetric.py

@@ -168,7 +168,7 @@ def in_manifold_eigen(self, L, eps=1e-6):
         if L.size(-1) > self.rank:
             # We compute the \infty-norm of the remaining dimension
             D = L[..., : -self.rank]
-            infty_norm_err = D.abs().max(dim=-1).values
+            infty_norm_err = D.abs().amax(dim=-1)
             if (infty_norm_err > 5.0 * eps).any():
                 return False
         return (L[..., -self.rank :] >= -eps).all().item()

test/test_lowrank.py

@@ -35,3 +35,12 @@ 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_fixed_rank_sample_clamps_singular_values(self):
+        manifold = FixedRank(size=(3, 3), rank=2).double()
+        _, singular_values, _ = manifold.sample(
+            init_=lambda X: X.zero_(), eps=0.25, factorized=True
+        )
+        self.assertTrue(
+            torch.equal(singular_values, torch.full_like(singular_values, 0.25))
+        )

test/test_positive_semidefinite.py

@@ -1,5 +1,7 @@
 from unittest import TestCase
 
+import torch
+
 from geotorch.pssdlowrank import PSSDLowRank
 from geotorch.pssdfixedrank import PSSDFixedRank
 from geotorch.pssd import PSSD
@@ -40,3 +42,10 @@ def test_positive_semidefinite_errors(self):
             PSD(size=(5, 2), f=3)
         with self.assertRaises(ValueError):
             PSD(size=(5, 3), f="fail")
+
+    def test_fixed_rank_sample_clamps_eigenvalues(self):
+        manifold = PSSDFixedRank(size=(3, 3), rank=2).double()
+        sample = manifold.sample(init_=lambda X: X.zero_(), eps=0.25)
+        eigenvalues = torch.linalg.eigvalsh(sample)
+        expected = torch.tensor([0.0, 0.25, 0.25], dtype=torch.float64)
+        torch.testing.assert_close(eigenvalues, expected)

test/test_sphere.py

@@ -51,3 +51,11 @@ def test_embedded_sample_uses_module_dtype(self):
         sample = sphere.sample()
         self.assertEqual(sample.dtype, torch.float64)
         self.assertEqual(sample.shape, (2, 3))
+
+    def test_membership_preserves_strict_epsilon_boundary(self):
+        sphere = SphereEmbedded(size=(1,), radius=2.0)
+        x = torch.tensor([2.0001], dtype=torch.float64)
+        eps = (torch.linalg.vector_norm(x) - sphere.radius).item()
+
+        self.assertFalse(sphere.in_manifold(x, eps=eps))
+        self.assertTrue(sphere.in_manifold(x, eps=2.0 * eps))