Compute multiple soft-quantiles in one execution without using `vmap` #382

marcocuturi · 2023-07-01T11:25:13Z

This PR is a follow up to
#373
This implements a quantiles function in the soft_sort module to return simultenaously multiple quantile values. Should be more efficient than the vmap proposed in the discussion, but will likely return very slightly different results.

marcocuturi · 2023-07-02T07:40:44Z

numba, via jaxopt, seems to be causing the issue.

review-notebook-app · 2023-07-03T10:28:20Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

codecov-commenter · 2023-07-03T12:45:05Z

Codecov Report

Merging #382 (4268043) into main (3b4f7b6) will increase coverage by 0.02%.
The diff coverage is 96.55%.

❗ Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please install the Github App Integration for your organization. Read more.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #382      +/-   ##
==========================================
+ Coverage   88.51%   88.53%   +0.02%     
==========================================
  Files          52       52              
  Lines        5660     5679      +19     
  Branches      839      841       +2     
==========================================
+ Hits         5010     5028      +18     
  Misses        530      530              
- Partials      120      121       +1

Impacted Files	Coverage Δ
src/ott/solvers/linear/sinkhorn.py	`96.94% <ø> (ø)`
src/ott/tools/soft_sort.py	`95.23% <96.55%> (-0.09%)`	⬇️

tests/tools/soft_sort_test.py

michalk8 · 2023-07-03T19:53:09Z

tests/tools/soft_sort_test.py

+    q = jnp.array([.1, .8, .4])
+    m1 = soft_sort.quantile(inputs, q=q, weight=None, axis=0)
+    np.testing.assert_allclose(m1.mean(axis=[1, 2]), q, atol=5e-2)
+    m2 = soft_sort.quantile(inputs, q=q, weight=None, axis=0)


Why does m2 exist since it's the same as m1?

thanks! I think this was naive byproduct of starting PR with 2 different interfaces (quantile and quantiles)

michalk8 · 2023-07-03T19:54:55Z

src/ott/tools/soft_sort.py

@@ -141,28 +146,50 @@ def sort(
 ) -> jnp.ndarray:
  r"""Apply the soft sort operator on a given axis of the input.

+  For instance:
+
+  ```


Use the python code blocks (also in other places), e.g.,:

.. code-block:: python x = jax.random.uniform(rng, (100,))

Otherwise, it doesn't render nicely, see here

michalk8 · 2023-07-03T19:58:07Z

src/ott/tools/soft_sort.py

+        jnp.ones((num_quantiles + 1, 1), dtype=bool)
+    ],
+                           axis=1).ravel()[:-1]
+    return (out[odds])[idx]


Unnecessary ()

michalk8 · 2023-07-03T20:00:09Z

src/ott/tools/soft_sort.py


  Returns:
-    A jnp.ndarray of the same shape as the input with soft sorted values on the
+    A jnp.ndarray of the same shape as the input with soft-sorted values on the


Maybe it's a good time to refactor and say An array of the same shape ...

michalk8 · 2023-07-03T20:04:10Z

src/ott/tools/soft_sort.py

+    targets: sorted array (in ascending order) of dimension 1 describing a
+      discrete distribution. Note: the``targets`` values must be provided as
+      a sorted vector.
+    weights: vector of nonnegative weights, summing to :math:`1.0`, of the same


non-negative

there's an argument for dropping the hyphen :) https://math.stackexchange.com/questions/3342643/nonnegative-vs-non-negative

michalk8 · 2023-07-03T20:04:24Z

src/ott/tools/soft_sort.py

+    targets: sorted array (in ascending order) of dimension 1 describing a
+      discrete distribution. Note: the``targets`` values must be provided as
+      a sorted vector.
+    weights: vector of nonnegative weights, summing to :math:`1.0`, of the same


michalk8 · 2023-07-03T20:04:52Z

src/ott/tools/soft_sort.py

+    inputs: array of any shape whose values will be changed to match those in
+      ``targets``.
+    targets: sorted array (in ascending order) of dimension 1 describing a
+      discrete distribution. Note: the``targets`` values must be provided as


the ``targets`` (missing space, doesn't render correctly)

michalk8 · 2023-07-03T20:05:44Z

src/ott/tools/soft_sort.py

-      specified by the optimal transport between values in ``inputs`` towards
-      those values. If not specified, ``num_targets`` is set by default to be
-      the size of the slices of the input that are sorted.
+    inputs: jnp.ndarray<float> of any shape.


Array of any shape.

Also elsewhere, if possible.

michalk8 · 2023-07-03T20:07:08Z

src/ott/tools/soft_sort.py

    num_points = inputs.shape[0]
+    q = jnp.array([0.2, 0.5, 0.8]) if q is None else jnp.atleast_1d(q)


Why are defaults needed?

Also, pure implementation question: why pass these as an array, not e.g., as a tuple? To be able to differentiate w.r.t. to it?

I think it's also a question of being able to re-run by changing quantile values, without jitting again, pending that the number of quantiles does not change.

you prefer no defaults? it was there because there was a default previously (the median) but maybe you're right, better stick to jax's quantile API

* comments by Michal in #382 * test

…#382) * add quantiles * numpy * using Michal's fix * chg threshold in kmeans test * changing quantile API to match jnp's + pydocs * impact chg in NB * pydocs

* comments by Michal in #382 * test

add quantiles

8cfe5af

marcocuturi changed the title ~~Possibility to compute multilple quantiles at once~~ Possibility to compute multiple quantiles at once Jul 1, 2023

marcocuturi changed the title ~~Possibility to compute multiple quantiles at once~~ Compute multiple soft-quantiles in one execution without using vmap Jul 1, 2023

numpy

1ddcfe3

marcocuturi added 4 commits July 2, 2023 14:58

using Michal's fix

872e3af

chg threshold in kmeans test

6a95cc4

changing quantile API to match jnp's + pydocs

66b4b12

impact chg in NB

422afe4

pydocs

4268043

marcocuturi requested a review from michalk8 July 3, 2023 14:23

marcocuturi merged commit b2b7ebb into main Jul 3, 2023
10 checks passed

marcocuturi deleted the mott4 branch July 3, 2023 21:48

michalk8 requested changes Jul 3, 2023

View reviewed changes

marcocuturi added a commit that referenced this pull request Jul 4, 2023

comments by Michal in #382

202ed76

marcocuturi added a commit that referenced this pull request Jul 4, 2023

incorporate feedback in #382 (#385)

57fb327

* comments by Michal in #382 * test

michalk8 mentioned this pull request Aug 1, 2023

Cant jit the function using soft_sort.quantile() and jnp.where() #407

Closed

michalk8 pushed a commit that referenced this pull request Jun 27, 2024

incorporate feedback in #382 (#385)

7fcc45f

* comments by Michal in #382 * test

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Compute multiple soft-quantiles in one execution without using `vmap` #382

Compute multiple soft-quantiles in one execution without using `vmap` #382

marcocuturi commented Jul 1, 2023

marcocuturi commented Jul 2, 2023

review-notebook-app bot commented Jul 3, 2023

codecov-commenter commented Jul 3, 2023 •

edited

Loading

michalk8 Jul 3, 2023

marcocuturi Jul 4, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

marcocuturi Jul 4, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

marcocuturi Jul 4, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

michalk8 Jul 3, 2023

marcocuturi Jul 4, 2023

marcocuturi Jul 4, 2023 •

edited

Loading

		num_points = inputs.shape[0]
		q = jnp.array([0.2, 0.5, 0.8]) if q is None else jnp.atleast_1d(q)

Compute multiple soft-quantiles in one execution without using vmap #382

Compute multiple soft-quantiles in one execution without using vmap #382

Conversation

marcocuturi commented Jul 1, 2023

marcocuturi commented Jul 2, 2023

review-notebook-app bot commented Jul 3, 2023

codecov-commenter commented Jul 3, 2023 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

marcocuturi Jul 4, 2023 • edited Loading

Choose a reason for hiding this comment

Compute multiple soft-quantiles in one execution without using `vmap` #382

Compute multiple soft-quantiles in one execution without using `vmap` #382

codecov-commenter commented Jul 3, 2023 •

edited

Loading

marcocuturi Jul 4, 2023 •

edited

Loading