restructuring neural models + addition of OT-FM and GENOT

.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.1.6
   hooks:
   - id: ruff
     args: [--fix, --exit-non-zero-on-fix]
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  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.11.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]

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"
 

docs/references.bib

@@ -814,3 +814,53 @@ @misc{huguet:2023
   title       = {Geodesic Sinkhorn for Fast and Accurate Optimal Transport on Manifolds},
   year        = {2023},
 }
+
+@misc{eyring:23,
+  author={Eyring, Luca and Klein, Dominik and Uscidda, Th{\'e}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={Dominik Klein and Théo Uscidda and Fabian Theis and Marco Cuturi},
+      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}
+}

docs/tutorials/MetaOT.ipynb

@@ -81,6 +81,8 @@
    "outputs": [],
    "source": [
     "# Obtain the MNIST dataset and flatten the images into discrete measures.\n",
+    "\n",
+    "\n",
     "def get_mnist_flat(train):\n",
     "    dataset = torchvision.datasets.MNIST(\n",
     "        \"/tmp/mnist/\",\n",

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",

docs/tutorials/icnn_inits.ipynb

@@ -33,6 +33,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",

docs/tutorials/neural_dual.ipynb

@@ -49,9 +49,10 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
-    "import optax\n",
     "from torch.utils.data import DataLoader, IterableDataset\n",
     "\n",
+    "import optax\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "from IPython.display import clear_output, display\n",
     "\n",

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",

docs/tutorials/soft_sort.ipynb

@@ -37,16 +37,17 @@
     "\n",
     "from tqdm.notebook import tqdm\n",
     "\n",
-    "import flax.linen as nn\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
-    "import optax\n",
     "import torchvision\n",
-    "from flax import struct\n",
     "from scipy import ndimage\n",
     "from torch.utils import data\n",
     "\n",
+    "import flax.linen as nn\n",
+    "import optax\n",
+    "from flax import struct\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
     "from ott.tools import soft_sort"

docs/tutorials/sparse_monge_displacements.ipynb

@@ -114,6 +114,8 @@
    "outputs": [],
    "source": [
     "# Plotting utility\n",
+    "\n",
+    "\n",
     "def plot_map(x, y, x_new=None, z=None, ax=None, title=None):\n",
     "    if ax is None:\n",
     "        f, ax = plt.subplots(figsize=(10, 8))\n",

docs/tutorials/tracking_progress.ipynb

@@ -373,6 +373,8 @@
    "outputs": [],
    "source": [
     "# Samples spiral\n",
+    "\n",
+    "\n",
     "def sample_spiral(\n",
     "    n, min_radius, max_radius, key, min_angle=0, max_angle=10, noise=1.0\n",
     "):\n",

pyproject.toml

@@ -59,6 +59,7 @@ Changelog = "https://github.com/ott-jax/ott/releases"
 neural = [
     "flax>=0.6.6",
     "optax>=0.1.1",
+    "diffrax>=0.4.1",
 ]
 dev = [
     "pre-commit>=2.16.0",
@@ -102,11 +103,14 @@ include = '\.ipynb$'
 
 [tool.isort]
 profile = "black"
+line_length = 80
 include_trailing_comma = true
 multi_line_output = 3
-sections = ["FUTURE", "STDLIB", "THIRDPARTY", "NUMERIC", "PLOTTING", "FIRSTPARTY", "LOCALFOLDER"]
-# also contains what we import in notebooks
-known_numeric = ["numpy", "scipy", "jax", "flax", "optax", "jaxopt", "torch", "ot", "torchvision", "pandas", "sklearn"]
+sections = ["FUTURE", "STDLIB", "THIRDPARTY", "TEST", "NUMERIC", "NEURAL", "PLOTTING", "FIRSTPARTY", "LOCALFOLDER"]
+# also contains what we import in notebooks/tests
+known_neural = ["flax", "optax", "diffrax", "orbax"]
+known_numeric = ["numpy", "scipy", "jax", "flax", "optax", "jaxopt", "torch", "ot", "torchvision", "pandas", "sklearn", "tslearn"]
+known_test = ["pytest"]
 known_plotting = ["IPython", "matplotlib", "mpl_toolkits", "seaborn"]
 
 [tool.pytest.ini_options]
@@ -285,7 +289,6 @@ ignore = [
 line-length = 80
 select = [
     "D", # flake8-docstrings
-    "I", # isort
     "E", # pycodestyle
     "F", # pyflakes
     "W", # pycodestyle
@@ -301,7 +304,7 @@ select = [
     "T20",  # flake8-print
     "RET", # flake8-raise
 ]
-unfixable = ["B", "UP", "C4", "BLE", "T20", "RET"]
+unfixable = ["I", "B", "UP", "C4", "BLE", "T20", "RET"]
 target-version = "py38"
 [tool.ruff.per-file-ignores]
 # TODO(michalk8): PO004 - remove `self.initialize`

src/ott/datasets.py

@@ -51,13 +51,13 @@ class GaussianMixture:
         rectangle
 
     batch_size: batch size of the samples
-    init_rng: initial PRNG key
+    rng: initial PRNG key
     scale: scale of the Gaussian means
     std: the standard deviation of the individual Gaussian samples
   """
   name: Name_t
   batch_size: int
-  init_rng: jax.Array
+  rng: jax.Array
   scale: float = 5.0
   std: float = 0.5
 
@@ -96,7 +96,7 @@ def __iter__(self) -> Iterator[jnp.array]:
     return self._create_sample_generators()
 
   def _create_sample_generators(self) -> Iterator[jnp.array]:
-    rng = self.init_rng
+    rng = self.rng
     while True:
       rng1, rng2, rng = jax.random.split(rng, 3)
       means = jax.random.choice(rng1, self.centers, (self.batch_size,))
@@ -128,26 +128,18 @@ def create_gaussian_mixture_samplers(
   rng1, rng2, rng3, rng4 = jax.random.split(rng, 4)
   train_dataset = Dataset(
       source_iter=iter(
-          GaussianMixture(
-              name_source, batch_size=train_batch_size, init_rng=rng1
-          )
+          GaussianMixture(name_source, batch_size=train_batch_size, rng=rng1)
       ),
       target_iter=iter(
-          GaussianMixture(
-              name_target, batch_size=train_batch_size, init_rng=rng2
-          )
+          GaussianMixture(name_target, batch_size=train_batch_size, rng=rng2)
       )
   )
   valid_dataset = Dataset(
       source_iter=iter(
-          GaussianMixture(
-              name_source, batch_size=valid_batch_size, init_rng=rng3
-          )
+          GaussianMixture(name_source, batch_size=valid_batch_size, rng=rng3)
       ),
       target_iter=iter(
-          GaussianMixture(
-              name_target, batch_size=valid_batch_size, init_rng=rng4
-          )
+          GaussianMixture(name_target, batch_size=valid_batch_size, rng=rng4)
       )
   )
   dim_data = 2

src/ott/math/__init__.py

@@ -11,9 +11,4 @@
 # 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 (
-    fixed_point_loop,
-    matrix_square_root,
-    unbalanced_functions,
-    utils,
-)
+from . import fixed_point_loop, matrix_square_root, unbalanced_functions, utils

src/ott/neural/__init__.py

@@ -11,4 +11,4 @@
 # 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 layers, losses, models, solvers
+from . import data, duality, flows, gaps, models

src/ott/neural/solvers/__init__.py → src/ott/neural/data/__init__.py

@@ -11,4 +11,4 @@
 # 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 conjugate, map_estimator, neuraldual
+from . import dataloaders