Repository navigation
Conversation
…directly nx.sqrtm took the square root of eigh's eigenvalues as returned in the NumPy, JAX, CuPy and TensorFlow backends. For a positive semi-definite matrix with zero eigenvalues, rounding can return them slightly negative, so the square root came back NaN. ot.gaussian.bures_distance then returned NaN on rank-deficient covariances, or raised LinAlgError when a NaN reached a second eigh. The torch backend already clamped to zero, so the same call gave a finite answer under torch and NaN under NumPy. All backends now clamp, and the base docstring says so. Rank-deficient covariances are the usual case when there are fewer samples than dimensions. Verified on numpy, jax, torch and tensorflow: B(C, C) is 0 at every rank, and the commuting closed form sqrt(sum (sqrt(a) - sqrt(b))^2) is matched with a shared null space. CuPy is changed identically but could not be run here. Clamping alone would have regressed bures_wasserstein_mapping. The map needs inv(sqrtm(Cs)), which does not exist for a singular source covariance. On master such inputs raised "array must not contain infs or NaNs" from scipy.linalg.inv (60 of 60 seeds, 30 samples in 50 dimensions); with only the clamp they returned a finite, wrong map without any warning (60 of 60, residual at least 3.9e33). The mapping now checks the identity that defines it, A Cs A = Ct, and warns when the relative residual exceeds sqrt(eps) of the dtype. This costs 1.2 to 1.5% of the mapping from d = 100 up. On well-posed and moderately ill-conditioned data with the default reg it does not fire (residuals 4e-15 to 5e-13); singular sources trip it every time. It also removes a false alarm. empirical_bures_wasserstein_mapping warned and returned NaN on 766 samples in 800 dimensions with reg=1e-8, a case where the source covariance is invertible and the map exists: with the clamp it is accurate to a relative residual of 9.8e-9. The test that pinned the NaN now asserts the warning on reg=0, where the map genuinely does not exist, and a new test asserts accuracy at reg=1e-8. bures_wasserstein_mapping is split into a public wrapper over a private body so that the warning points at the user's call site whether it comes from the public function or the empirical one.
# Conflicts: # RELEASES.md
Correct maps between rank-deficient covariances reach residuals near sqrt(eps): 1.0e-8 at the default reg on OpenBLAS and 1.1e-8 at reg=1e-8 on Accelerate (766 samples in 800 dimensions, worst of five seeds), against a tolerance of 1.5e-8. eps**0.25 is 1.2e-4 in float64. It still warns on singular sources (residual above 1e37), from reg=1e-12 (3e-4 to 3e-3), and on float32 maps of the same data at the default reg, which are off by 7 to 13%.
jax.grad through bures_wasserstein_mapping raised TracerArrayConversionError: the check read traced values back to NumPy outside the try meant to skip them. jax.jit only worked because the dtype probe failed first. All reads are now inside the try. The new test differentiates with respect to each covariance separately; with respect to both, the probe fails first and hides the bug.
Also reword the entry: the non-finite check in empirical_bures_wasserstein_mapping is still there.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #905 +/- ##
==========================================
- Coverage 97.01% 97.01% -0.01%
==========================================
Files 130 130
Lines 26879 26951 +72
==========================================
+ Hits 26078 26146 +68
- Misses 801 805 +4 🚀 New features to boost your workflow:
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Types of changes
Bug fix (non-breaking change which fixes an issue).
Motivation and context / Related issue
Closes #889.
nx.sqrtmnow clamps negative eigenvalues to zero in all backends, as torch already did (#774). CuPy untested.Clamping alone would turn the
scipy.linalg.inverror inbures_wasserstein_mappingfor a singular source into a silently wrong map (residual ≥ 3.9e33, 60 of 60 seeds). It now warns when ‖A Cs A − Ct‖ / ‖Ct‖ > ε^¼ (1.2e-4 in float64). √ε (1.5e-8) is too close to what correct maps reach on rank-deficient data (1.1e-8). Overhead 1.2–1.5% from d=100.This also fixes NaN from
empirical_bures_wasserstein_mappingatreg=1e-8(766 samples, 800 dimensions); the test pinning it now checks the warning atreg=0.For a singular source this warns and returns, like
empirical_bures_wasserstein_mapping; happy to raise instead.How has this been tested (if it applies)
New and changed tests fail on master (except torch) and pass here; full suite passes.
PR checklist