[MRG] Stabilize non-Torch matrix square roots for rank-deficient covariances

RELEASES.md

@@ -18,6 +18,7 @@
 
 #### Closed issues
 
+- Clamp roundoff-negative eigenvalues in non-Torch backend matrix square roots so rank-deficient positive semi-definite covariances remain finite (Issue #889)
 - Fix duplicate function and class descriptions in generated module API pages while preserving top-level aliases and example backreferences (Issue #527)
 - Fix internal arrays created on the torch default device instead of the input device in `ot.partial.partial_wasserstein`, `ot.wasserstein_circle` (and `ot.sliced_wasserstein_sphere`), the high-dimensional Bures-Wasserstein functions in `ot.gaussian`, `ot.gmm`, `ot.lowrank_sinkhorn` with `init="deterministic"`, the semi-relaxed (F)GW barycenters and `ot.utils.projection_sparse_simplex`, which made them fail on GPU inputs. Add a `meta_default_device` test fixture to catch such allocations on CPU (PR #883, Issue #852)
 - Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848)

ot/backend.py

@@ -1574,7 +1574,7 @@ def pinv(self, a, hermitian=False):
 
     def sqrtm(self, a):
         L, V = np.linalg.eigh(a)
-        L = np.sqrt(L)
+        L = np.sqrt(np.maximum(L, 0))
         # Q[...] = V[...] @ diag(L[...])
         Q = np.einsum("...jk,...k->...jk", V, L)
         # R[...] = Q[...] @ V[...].T
@@ -2027,7 +2027,7 @@ def pinv(self, a, hermitian=False):
 
     def sqrtm(self, a):
         L, V = jnp.linalg.eigh(a)
-        L = jnp.sqrt(L)
+        L = jnp.sqrt(jnp.maximum(L, 0))
         # Q[...] = V[...] @ diag(L[...])
         Q = jnp.einsum("...jk,...k->...jk", V, L)
         # R[...] = Q[...] @ V[...].T
@@ -3107,7 +3107,7 @@ def pinv(self, a, hermitian=False):
 
     def sqrtm(self, a):
         L, V = cp.linalg.eigh(a)
-        L = cp.sqrt(L)
+        L = cp.sqrt(cp.maximum(L, 0))
         # Q[...] = V[...] @ diag(L[...])
         Q = cp.einsum("...jk,...k->...jk", V, L)
         # R[...] = Q[...] @ V[...].T
@@ -3586,7 +3586,7 @@ def pinv(self, a, hermitian=False):
 
     def sqrtm(self, a):
         L, V = tf.linalg.eigh(a)
-        L = tf.sqrt(L)
+        L = tf.sqrt(tf.maximum(L, 0))
         # Q[...] = V[...] @ diag(L[...])
         Q = tf.einsum("...jk,...k->...jk", V, L)
         # R[...] = Q[...] @ V[...].T

test/test_backend.py

@@ -812,6 +812,20 @@ def test_func_backends(nx):
         )
 
 
+def test_sqrtm_rank_deficient_is_finite(nx):
+    rng = np.random.RandomState(0)
+    factors = rng.randn(10, 5)
+    covariance = factors @ factors.T / 5
+
+    covariance_backend = nx.from_numpy(covariance)
+    square_root = nx.to_numpy(nx.sqrtm(covariance_backend))
+
+    assert np.isfinite(square_root).all()
+    np.testing.assert_allclose(
+        square_root @ square_root, covariance, rtol=1e-5, atol=1e-6
+    )
+
+
 def test_random_backends(nx):
     tmp_u = nx.rand()
 

test/test_gaussian.py

@@ -156,13 +156,12 @@ def test_empirical_bures_wasserstein_mapping_hd(nx, bias):
     np.testing.assert_allclose(Ct, Cst, rtol=1.0, atol=1.0)
 
 
-def test_empirical_bures_wasserstein_mapping_numerical_error_warning():
+def test_empirical_bures_wasserstein_mapping_ill_conditioned_is_finite():
     rng = np.random.RandomState(42)
     Xs = rng.rand(766, 800) * 5
     Xt = rng.rand(295, 800) * 2
-    with pytest.warns():
-        A, b = ot.gaussian.empirical_bures_wasserstein_mapping(Xs, Xt, reg=1e-8)
-        assert not is_all_finite(A, b)
+    A, b = ot.gaussian.empirical_bures_wasserstein_mapping(Xs, Xt, reg=1e-8)
+    assert is_all_finite(A, b)
 
 
 def test_bures_wasserstein_distance(nx):