Compute multiple soft-quantiles in one execution without using vmap (…

…#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

docs/tutorials/notebooks/soft_sort.ipynb

@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "Dg9-8gVqjq2H"
@@ -172,7 +173,7 @@
     }
    ],
    "source": [
-    "jnp.quantile(x, 0.5)"
+    "jnp.quantile(x, q=0.5)"
    ]
   },
   {
@@ -196,6 +197,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "mnEkXjwT-Z1C"
@@ -334,6 +336,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "JKNcCJOe9Dcl"
@@ -377,6 +380,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "4w4ggUy7zYQX"
@@ -442,6 +446,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "jQoRgYEd-pkj"
@@ -494,7 +499,7 @@
    ],
    "source": [
     "softquantile = jax.jit(soft_sort.quantile)\n",
-    "softquantile(x, level=0.5)"
+    "softquantile(x, q=0.5)"
    ]
   },
   {
@@ -533,6 +538,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -596,6 +602,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "4t3VrtNcmN0R"
@@ -611,6 +618,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "tqsCC0tunHQh"
@@ -663,6 +671,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "irSGHYZ7nWuY"
@@ -709,6 +718,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "jK94muT8oAlQ"
@@ -929,6 +939,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -976,6 +987,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [

pyproject.toml

@@ -16,8 +16,8 @@ dependencies = [
     "jax>=0.1.67",
     "jaxlib>=0.1.47",
     "jaxopt>=0.5.5",
-    # https://github.com/google/jax/discussions/9951#discussioncomment-3017784
-    "numpy>=1.18.4, !=1.23.0",
+    ## https://github.com/google/jax/discussions/9951#discussioncomment-3017784
+    "numpy>=1.18.4, !=1.25.0",
     "flax>=0.5.2",
     "optax>=0.1.1",
     "lineax>=0.0.1; python_version >= '3.9'"

src/ott/solvers/linear/sinkhorn.py

@@ -732,7 +732,9 @@ class Sinkhorn:
       gradients have been stopped. This is useful when carrying out first order
       differentiation, and is only valid (as with ``implicit_differentiation``)
       when the algorithm has converged with a low tolerance.
-    initializer: how to compute the initial potentials/scalings.
+    initializer: how to compute the initial potentials/scalings. This refers to
+      a few possible classes implemented following the template in
+      :class:`~ott.initializers.linear.SinkhornInitializer`.
     progress_fn: callback function which gets called during the Sinkhorn
       iterations, so the user can display the error at each iteration,
       e.g., using a progress bar. See :func:`~ott.utils.default_progress_fn`