ott-jax · michalk8 · Jun 21, 2024 · Jun 18, 2024 · Jun 18, 2024 · Jun 18, 2024
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -57,7 +57,7 @@ jobs:
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install -e".[test]"
-        python3 -m pip install "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+        python3 -m pip install "jax[cuda12]"
 
     - name: Run nvidia-smi
       run: |

diff --git a/docs/experimental/index.rst b/docs/experimental/index.rst
@@ -0,0 +1,11 @@
+ott.experimental
+================
+.. module:: ott.experimental
+
+The :mod:`ott.experimental` module groups experimental code that might be useful
+for users but whose API we expect to change significantly in coming months.
+
+.. toctree::
+    :maxdepth: 2
+
+    mmsinkhorn
diff --git a/docs/experimental/mmsinkhorn.rst b/docs/experimental/mmsinkhorn.rst
@@ -0,0 +1,15 @@
+ott.experimental.mmsinkhorn
+===========================
+.. module:: ott.experimental.mmsinkhorn
+.. currentmodule:: ott.experimental.mmsinkhorn
+
+Solvers for multimarginal entropic OT problems, defined using :math:`k` point
+clouds of variable sizes in dimension :math:`d`.
+
+Sinkhorn Initializers
+---------------------
+.. autosummary::
+    :toctree: _autosummary
+
+    MMSinkhorn
+    MMSinkhornOutput
diff --git a/docs/index.rst b/docs/index.rst
@@ -117,6 +117,7 @@ Packages
     solvers/index
     initializers/index
     neural/index
+    experimental/index
     tools
     math
     utils

diff --git a/docs/references.bib b/docs/references.bib
@@ -25,6 +25,19 @@ @article{cuturi:15
   year    = {2016},
 }
 
+
+@inproceedings{piran:24,
+  author    = {Piran, Zoe and Klein, Michal and Thornton, James and Cuturi, Marco},
+  publisher = {PMLR},
+  booktitle = {Proceedings of The 41st International Conference on Machine Learning},
+  month     = jun,
+  series    = {Proceedings of Machine Learning Research},
+  title     = {Contrasting Multiple Representations with the Multi-Marginal Matching Gap},
+  year      = {2024},
+}
+
+
+
 @inproceedings{peyre:16,
   author    = {Peyré, Gabriel and Cuturi, Marco and Solomon, Justin},
   editor    = {Balcan, Maria Florina and Weinberger, Kilian Q.},
@@ -356,6 +369,19 @@ @article{solomon:15
   year      = {2015},
 }
 
+
+@article{gangbo:98,
+  author = {Gangbo, Wilfrid and Święch, Andrzej},
+  title = {Optimal maps for the multidimensional Monge-Kantorovich problem},
+  journal = {Communications on Pure and Applied Mathematics},
+  volume = {51},
+  number = {1},
+  pages = {23-45},
+  year = {1998}
+}
+
+
+
 @inproceedings{genevay:18,
   author    = {Genevay, Aude and Peyre, Gabriel and Cuturi, Marco},
   editor    = {Storkey, Amos and Perez-Cruz, Fernando},

diff --git a/docs/spelling/technical.txt b/docs/spelling/technical.txt
@@ -88,6 +88,7 @@ linearized
 logit
 macOS
 methylation
+multimarginal
 neuroimaging
 normed
 numerics
@@ -100,6 +101,7 @@ parameterization
 parameterizing
 piecewise
 pluripotent
+polymatching
 polynomials
 positivity
 postfix

diff --git a/docs/tutorials/Sinkhorn_Barycenters.ipynb b/docs/tutorials/Sinkhorn_Barycenters.ipynb
@@ -6,7 +6,9 @@
     "id": "C4nZjbMHWcm_"
    },
    "source": [
-    "# Sinkhorn Barycenters"
+    "# Sinkhorn Barycenters\n",
+    "\n",
+    "This tutorial covers the computation of Wasserstein barycenters using regularized OT, based on {cite}`cuturi:14` and {cite}`benamou:15`."
    ]
   },
   {
@@ -56,20 +58,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {
     "id": "ysURew0UKhHE"
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/michal/mambaforge/envs/ott/lib/python3.10/site-packages/nilearn/datasets/struct.py:850: UserWarning: `legacy_format` will default to `False` in release 0.11. Dataset fetchers will then return pandas dataframes by default instead of recarrays.\n",
-      "  warnings.warn(_LEGACY_FORMAT_MSG)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "n_subjects = 4\n",
     "dataset_files = datasets.fetch_oasis_vbm(\n",

diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -18,6 +18,7 @@ Linear Optimal Transport
     OTT_&_POT
     Hessians
     LRSinkhorn
+    mmsink
     sinkhorn_divergence_gradient_flow
     sparse_monge_displacements
 

diff --git a/docs/tutorials/mmsink.ipynb b/docs/tutorials/mmsink.ipynb
diff --git a/src/ott/__init__.py b/src/ott/__init__.py
@@ -15,6 +15,7 @@
 
 from . import (
     datasets,
+    experimental,
     geometry,
     initializers,
     math,

diff --git a/src/ott/experimental/__init__.py b/src/ott/experimental/__init__.py
@@ -0,0 +1,14 @@
+# Copyright OTT-JAX
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from . import mmsinkhorn