[gezerui888]

Types of changes

Bug fix.

Motivation and context / Related issue

Addresses #625. This fixes the emd2 loss path discussed there; it does not change the transport-plan API or add native jit/vmap support.

jax.grad(ot.emd2) fails when traced inputs reach the native solver. I detach the solver inputs and keep the original arrays for gradient registration. I also replace the JAX gradient attachment so that it uses the supplied gradients at full magnitude and preserves their dependence for higher derivatives. This also corrects the JAX envelope gradients returned by ot.solve.

How has this been tested (if it applies)

I added 36 cases for marginal and cost gradients, finite differences, empty marginals, dtype and return options, upstream scaling, nested values, and higher derivatives.

• pytest --durations=20 -v test/ --doctest-modules: 3,592 passed, 487 skipped, six xfailed. Base has 3,556 passed with the same skips and xfails.
• pre-commit run --all-files: passed.

PR checklist

• I have read the CONTRIBUTING document.
• The documentation is up-to-date with the changes I made (check build artifacts).
• All tests passed, and additional code has been covered with new tests.
• I have added the PR and Issue fix to the RELEASES.md file.

cc @rflamary

[gezerui888]

@rflamary I'd appreciate a review when you have time. Could you also approve the pending CI workflows?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.09%. Comparing base (3367480) to head (5c97fc4).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #917      +/-   ##
==========================================
+ Coverage   97.07%   97.09%   +0.01%     
==========================================
  Files         130      130              
  Lines       27095    27229     +134     
==========================================
+ Hits        26302    26437     +135     
+ Misses        793      792       -1     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.