introduce new norm operator that can differentiate through lambda x: norm(x-x) + lighter tests across the board
#411
Conversation
This file contains 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
marcocuturi
changed the title
norm operator that can differentiate through lambda x: norm(x-x)
|
also of interest to @theouscidda6 as potentially useful in Monge gap. |
michalk8
requested changes
Aug 14, 2023
| # Test div of x to itself close to 0. | ||
| # Check differentiability of Sinkhorn divergence works, without NaN's. | ||
| grad = jax.grad(lambda x: div(x).divergence)(x) | ||
| assert jnp.all(jnp.logical_not(jnp.isnan(grad))) |
There was a problem hiding this comment.
not sure I can find the adequate test in there, https://numpy.org/doc/stable/reference/routines.testing.html
There was a problem hiding this comment.
You can use np.testing.assert_array_equal(jnp.isnan(grad), False).
|
thanks Michal! |
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #411 +/- ##
==========================================
- Coverage 91.51% 91.44% -0.07%
==========================================
Files 54 54
Lines 5891 5902 +11
Branches 856 860 +4
==========================================
+ Hits 5391 5397 +6
- Misses 364 370 +6
+ Partials 136 135 -1
|
marcocuturi
changed the title
norm operator that can differentiate through lambda x: norm(x-x)norm operator that can differentiate through lambda x: norm(x-x) + lighter tests across the board
michalk8
requested changes
Aug 15, 2023
| # Test div of x to itself close to 0. | ||
| # Check differentiability of Sinkhorn divergence works, without NaN's. | ||
| grad = jax.grad(lambda x: div(x).divergence)(x) | ||
| assert jnp.all(jnp.logical_not(jnp.isnan(grad))) |
There was a problem hiding this comment.
You can use np.testing.assert_array_equal(jnp.isnan(grad), False).
michalk8
pushed a commit
that referenced
this pull request
Jun 27, 2024
…x: norm(x-x)` + lighter tests across the board (#411) * introduce new norm operator * add * resolve name conflict in (math).utils_test.py * docs * vjp -> jvp * fix * add axis in test * type error * speed up some tests * fix * fix * fix apply_jacobian test * fix mem size * batch size * yet another mem fix * minor fixes + add in docs * minor docs fixes * docs
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.
Computing the gradient of Sinkhorn divergences using the
Euclideancost results inNaNvalues, because of the reliance onjnp.linalg.norm, and in particular, differentiation of the distance of a point against itself, e.g.norm(x-x)in the diagonal of a symmetric cost matrix.The fact that such a gradient is (and should be, in general)
NaNis well documented, see e.g. jax-ml/jax#6484However, In the context of OT, this poses problems, since it is safe to ignore these contributions, and therefore treat them as having 0 gradient.
This PR introduces a new
normfunction that does not blow up with aNaNat0.results in
as a result it should be now possible to differentiate through a
sinkhorn_divergencewith a more elaborate cost without producingNaN's.Also, in order to speed things up in CI, prune out some tests.