ott-jax · michalk8 · Apr 3, 2024 · Nov 22, 2023 · Nov 22, 2023 · Nov 22, 2023
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -6,25 +6,8 @@ default_stages:
 - push
 minimum_pre_commit_version: 3.0.0
 repos:
-- repo: https://github.com/google/yapf
-  rev: v0.40.0
-  hooks:
-  - id: yapf
-    additional_dependencies: [toml]
-- repo: https://github.com/nbQA-dev/nbQA
-  rev: 1.7.0
-  hooks:
-  - id: nbqa-pyupgrade
-    args: [--py38-plus]
-  - id: nbqa-black
-  - id: nbqa-isort
-- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-  rev: v2.10.0
-  hooks:
-  - id: pretty-format-yaml
-    args: [--autofix, --indent, '2']
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.4.0
+  rev: v4.5.0
   hooks:
   - id: detect-private-key
   - id: check-ast
@@ -37,13 +20,34 @@ repos:
   - id: trailing-whitespace
   - id: check-case-conflict
 - repo: https://github.com/charliermarsh/ruff-pre-commit
-  # Ruff version.
-  rev: v0.0.285
+  rev: v0.2.1
   hooks:
   - id: ruff
     args: [--fix, --exit-non-zero-on-fix]
+- repo: https://github.com/pycqa/isort
+  rev: 5.13.2
+  hooks:
+  - id: isort
+    name: isort
+- repo: https://github.com/google/yapf
+  rev: v0.40.2
+  hooks:
+  - id: yapf
+    additional_dependencies: [toml]
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.7.1
+  hooks:
+  - id: nbqa-pyupgrade
+    args: [--py38-plus]
+  - id: nbqa-black
+  - id: nbqa-isort
+- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
+  rev: v2.12.0
+  hooks:
+  - id: pretty-format-yaml
+    args: [--autofix, --indent, '2']
 - repo: https://github.com/rstcheck/rstcheck
-  rev: v6.1.2
+  rev: v6.2.0
   hooks:
   - id: rstcheck
     additional_dependencies: [tomli]

diff --git a/docs/conf.py b/docs/conf.py
@@ -26,9 +26,10 @@
 import logging
 from datetime import datetime
 
-import ott
 from sphinx.util import logging as sphinx_logging
 
+import ott
+
 # -- Project information -----------------------------------------------------
 needs_sphinx = "4.0"
 

diff --git a/docs/neural/data.rst b/docs/neural/data.rst
@@ -0,0 +1,15 @@
+ott.neural.data
+===============
+.. module:: ott.neural.data
+.. currentmodule:: ott.neural.data
+
+The :mod:`ott.problems.data` contains data sets and data loaders needed
+for solving (conditional) neural optimal transport problems.
+
+Datasets
+--------
+.. autosummary::
+    :toctree: _autosummary
+
+    datasets.OTDataset
+    datasets.ConditionalLoader
diff --git a/docs/neural/solvers.rst → docs/neural/duality.rst b/docs/neural/solvers.rst → docs/neural/duality.rst
@@ -1,20 +1,17 @@
-ott.neural.solvers
+ott.neural.duality
 ==================
-.. module:: ott.neural.solvers
-.. currentmodule:: ott.neural.solvers
+.. module:: ott.neural.duality
+.. currentmodule:: ott.neural.duality
 
 This module implements various solvers to estimate optimal transport between
 two probability measures, through samples, parameterized as neural networks.
-These neural networks are described in :mod:`ott.neural.models`, borrowing
-lower-level components from :mod:`ott.neural.layers` using
-`flax <https://flax.readthedocs.io/en/latest/>`__.
+These solvers build upon dual formulation of the optimal transport problem.
 
 Solvers
 -------
 .. autosummary::
     :toctree: _autosummary
 
-    map_estimator.MapEstimator
     neuraldual.W2NeuralDual
     neuraldual.BaseW2NeuralDual
 
@@ -26,3 +23,15 @@ Conjugate Solvers
     conjugate.FenchelConjugateLBFGS
     conjugate.FenchelConjugateSolver
     conjugate.ConjugateResults
+
+Models
+------
+.. autosummary::
+    :toctree: _autosummary
+
+    neuraldual.W2NeuralTrainState
+    neuraldual.BaseW2NeuralDual
+    neuraldual.W2NeuralDual
+    models.ICNN
+    models.PotentialMLP
+    models.MetaInitializer
diff --git a/docs/neural/flow_models.rst b/docs/neural/flow_models.rst
@@ -0,0 +1,48 @@
+ott.neural.flow_models
+======================
+.. module:: ott.neural.flow_models
+.. currentmodule:: ott.neural.flow_models
+
+This module implements various solvers building upon flow matching
+:cite:`lipman:22` to match distributions.
+
+Flows
+-----
+.. autosummary::
+    :toctree: _autosummary
+
+    flows.BaseFlow
+    flows.StraightFlow
+    flows.ConstantNoiseFlow
+    flows.BrownianNoiseFlow
+
+OT Flow Matching
+----------------
+.. autosummary::
+    :toctree: _autosummary
+
+    otfm.OTFlowMatching
+
+GENOT
+-----
+.. autosummary::
+    :toctree: _autosummary
+
+    genot.GENOTBase
+    genot.GENOTLin
+    genot.GENOTQuad
+
+Models
+------
+.. autosummary::
+    :toctree: _autosummary
+
+    models.VelocityField
+
+Utils
+-----
+.. autosummary::
+    :toctree: _autosummary
+
+    layers.CyclicalTimeEncoder
+    samplers.uniform_sampler
diff --git a/docs/neural/gaps.rst b/docs/neural/gaps.rst
@@ -0,0 +1,15 @@
+ott.neural.gaps
+===============
+.. module:: ott.neural.gaps
+.. currentmodule:: ott.neural.gaps
+
+This module implements gap models.
+
+Monge gap
+---------
+.. autosummary::
+    :toctree: _autosummary
+
+    map_estimator.MapEstimator
+    monge_gap.monge_gap
+    monge_gap.monge_gap_from_samples
diff --git a/docs/neural/index.rst b/docs/neural/index.rst
@@ -13,29 +13,8 @@ and solvers to estimate such neural networks.
 .. toctree::
     :maxdepth: 2
 
-    solvers
-
-Models
-------
-.. autosummary::
-    :toctree: _autosummary
-
-    models.ICNN
-    models.MLP
-    models.MetaInitializer
-
-Losses
-------
-.. autosummary::
-    :toctree: _autosummary
-
-    losses.monge_gap
-    losses.monge_gap_from_samples
-
-Layers
-------
-.. autosummary::
-    :toctree: _autosummary
-
-    layers.PositiveDense
-    layers.PosDefPotentials
+    data
+    duality
+    flow_models
+    gaps
+    models
diff --git a/docs/neural/models.rst b/docs/neural/models.rst
@@ -0,0 +1,25 @@
+ott.neural.models
+=================
+.. module:: ott.neural.models
+.. currentmodule:: ott.neural.models
+
+This module implements models, network architectures and helper
+functions which apply to various neural optimal transport solvers.
+
+Utils
+-----
+.. autosummary::
+    :toctree: _autosummary
+
+    base_solver.BaseOTMatcher
+    base_solver.OTMatcherLinear
+    base_solver.OTMatcherQuad
+
+
+Neural networks
+---------------
+.. autosummary::
+    :toctree: _autosummary
+
+    layers.MLPBlock
+    nets.RescalingMLP
diff --git a/docs/references.bib b/docs/references.bib
@@ -827,6 +827,56 @@ @misc{huguet:2023
   year        = {2023},
 }
 
+@misc{eyring:23,
+  author      = {Eyring, Luca and Klein, Dominik and Uscidda, Théo and Palla, Giovanni and Kilbertus, Niki and Akata, Zeynep and Theis, Fabian},
+  doi         = {10.48550/arXiv.2311.15100},
+  eprint      = {2311.15100},
+  eprintclass = {stat.ML},
+  eprinttype  = {arXiv},
+  title       = {Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation},
+  year        = {2023},
+}
+
+@misc{klein_uscidda:23,
+  author      = {Klein, Dominik and Uscidda, Théo and Theis, Fabian and Cuturi, Marco},
+  doi         = {10.48550/arXiv.2310.09254},
+  eprint      = {2310.09254},
+  eprintclass = {stat.ML},
+  eprinttype  = {arXiv},
+  title       = {Generative Entropic Neural Optimal Transport To Map Within and Across Spaces},
+  year        = {2023},
+}
+
+@misc{lipman:22,
+  author      = {Lipman, Yaron and Chen, Ricky TQ and Ben-Hamu, Heli and Nickel, Maximilian and Le, Matt},
+  doi         = {10.48550/arXiv.2210.02747},
+  eprint      = {2210.02747},
+  eprintclass = {stat.ML},
+  eprinttype  = {arXiv},
+  title       = {Flow matching for generative modeling},
+  year        = {2022},
+}
+
+@misc{tong:23,
+  author      = {Tong, Alexander and Malkin, Nikolay and Huguet, Guillaume and Zhang, Yanlei and {Rector-Brooks}, Jarrid and Fatras, Kilian and Wolf, Guy and Bengio, Yoshua},
+  doi         = {10.48550/arXiv.2302.00482},
+  eprint      = {2302.00482},
+  eprintclass = {stat.ML},
+  eprinttype  = {arXiv},
+  title       = {Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport},
+  year        = {2023},
+}
+
+@misc{pooladian:23,
+  author      = {Pooladian, Aram-Alexandre and Ben-Hamu, Heli and Domingo-Enrich, Carles and Amos, Brandon and Lipman, Yaron and Chen, Ricky},
+  doi         = {10.48550/arXiv.2304.14772},
+  eprint      = {2304.14772},
+  eprintclass = {stat.ML},
+  eprinttype  = {arXiv},
+  title       = {Multisample flow matching: Straightening flows with minibatch couplings},
+  year        = {2023},
+}
+
 @article{iacono:17,
   author  = {Iacono, Roberto and Boyd, John P.},
   url     = {https://doi.org/10.1007/s10444-017-9530-3},

diff --git a/docs/tutorials/MetaOT.ipynb b/docs/tutorials/MetaOT.ipynb
@@ -63,6 +63,7 @@
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
     "import torchvision\n",
+    "\n",
     "from flax import linen as nn\n",
     "\n",
     "import matplotlib.pyplot as plt\n",

diff --git a/docs/tutorials/Monge_Gap.ipynb b/docs/tutorials/Monge_Gap.ipynb
@@ -31,8 +31,9 @@
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "import optax\n",
     "import sklearn.datasets\n",
+    "\n",
+    "import optax\n",
     "from flax import linen as nn\n",
     "\n",
     "from matplotlib import pyplot as plt\n",

diff --git a/docs/tutorials/icnn_inits.ipynb b/docs/tutorials/icnn_inits.ipynb
@@ -32,6 +32,7 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
+    "\n",
     "import optax\n",
     "\n",
     "import matplotlib.pyplot as plt\n",

diff --git a/docs/tutorials/neural_dual.ipynb b/docs/tutorials/neural_dual.ipynb
@@ -37,8 +37,10 @@
    "source": [
     "import jax\n",
     "import jax.numpy as jnp\n",
+    "import numpy as np\n",
+    "from torch.utils.data import DataLoader, IterableDataset\n",
+    "\n",
     "import optax\n",
-    "from flax import linen as nn\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "from IPython.display import clear_output, display\n",

diff --git a/docs/tutorials/point_clouds.ipynb b/docs/tutorials/point_clouds.ipynb
@@ -64,7 +64,7 @@
    },
    "outputs": [],
    "source": [
-    "def create_points(rng: jax.random.PRNGKeyArray, n: int, m: int, d: int):\n",
+    "def create_points(rng: jax.Array, n: int, m: int, d: int):\n",
     "    rngs = jax.random.split(rng, 3)\n",
     "    x = jax.random.normal(rngs[0], (n, d)) + 1\n",
     "    y = jax.random.uniform(rngs[1], (m, d))\n",
@@ -279,6 +279,8 @@
    "outputs": [],
    "source": [
     "# Helper function to plot successively the optimal transports\n",
+    "\n",
+    "\n",
     "def plot_ots(ots):\n",
     "    fig = plt.figure(figsize=(8, 5))\n",
     "    plott = ott.tools.plot.Plot(fig=fig)\n",
@@ -366973,6 +366975,8 @@
    "outputs": [],
    "source": [
     "# Plotting utility\n",
+    "\n",
+    "\n",
     "def plot_map(x, y, z, forward: bool = True):\n",
     "    plt.figure(figsize=(10, 8))\n",
     "    marker_t = \"o\" if forward else \"X\"\n",