From 97d76149a0d32fb2ed4b8405b1dc5e22ea317aed Mon Sep 17 00:00:00 2001
From: Nikhila-Dhevi-R <dhevirnikhila@gmail.com>
Date: Fri, 23 Oct 2026 00:37:56 +0530
Subject: [PATCH] Fix backend sqrtm for rank-deficient covariances

---
 RELEASES.md           |  1 +
 ot/backend.py         |  8 ++++----
 test/test_backend.py  | 14 ++++++++++++++
 test/test_gaussian.py |  7 +++----
 4 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/RELEASES.md b/RELEASES.md
index d01ff927f..a62ae50c2 100644
--- a/RELEASES.md
+++ b/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)
diff --git a/ot/backend.py b/ot/backend.py
index f55938a93..c0a2c3139 100644
--- a/ot/backend.py
+++ b/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
diff --git a/test/test_backend.py b/test/test_backend.py
index e5b5ebd94..5582c017e 100644
--- a/test/test_backend.py
+++ b/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()
 
diff --git a/test/test_gaussian.py b/test/test_gaussian.py
index 3199e2403..f01bc40b4 100644
--- a/test/test_gaussian.py
+++ b/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):