Feature/batched vmap

docs/tutorials/barycenter/000_Sinkhorn_Barycenters.ipynb

@@ -17,9 +17,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import nilearn\n",
     "from nilearn import datasets, image, plotting\n",
-    "from nilearn.image import get_data\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
@@ -38,7 +36,7 @@
     "id": "02VJX2uXYHDX"
    },
    "source": [
-    "## Import neuroimaging data using `nilearn`.\n",
+    "## Import neuroimaging data using `nilearn`\n",
     "\n",
     "We recover a few MRI data points..."
    ]
@@ -157,7 +155,7 @@
     }
    ],
    "source": [
-    "a = jnp.array(get_data(gm_imgs)).transpose((3, 0, 1, 2))\n",
+    "a = jnp.array(image.get_data(gm_imgs)).transpose((3, 0, 1, 2))\n",
     "grid_size = a.shape[1:4]\n",
     "a = a.reshape((n_subjects, -1)) + 1e-2\n",
     "a = a / np.sum(a, axis=1)[:, np.newaxis]\n",
@@ -320,9 +318,7 @@
    ],
    "source": [
     "def data_to_nii(x):\n",
-    "    return nilearn.image.new_img_like(\n",
-    "        gm_imgs[0], data=np.array(x.reshape(grid_size))\n",
-    "    )\n",
+    "    return image.new_img_like(gm_imgs[0], data=np.array(x.reshape(grid_size)))\n",
     "\n",
     "\n",
     "plotting.plot_epi(data_to_nii(barycenter.histogram))\n",

docs/tutorials/barycenter/200_gmm_pair_demo.ipynb

@@ -69,7 +69,6 @@
     "    fit_gmm_pair,\n",
     "    gaussian_mixture,\n",
     "    gaussian_mixture_pair,\n",
-    "    probabilities,\n",
     ")"
    ]
   },

docs/tutorials/geometry/100_grid.ipynb

@@ -241,11 +241,15 @@
     "class MyCost(costs.CostFn):\n",
     "    \"\"\"An unusual cost function.\"\"\"\n",
     "\n",
-    "    def norm(self, x):\n",
+    "    def norm(self, x: jnp.ndarray) -> jnp.ndarray:\n",
     "        return jnp.sum(x**3 + jnp.cos(x) ** 2, axis=-1)\n",
     "\n",
-    "    def pairwise(self, x, y):\n",
-    "        return -jnp.sum(jnp.sin(x + 1) * jnp.sin(y)) * 2"
+    "    def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:\n",
+    "        return (\n",
+    "            self.norm(x)\n",
+    "            + self.norm(y)\n",
+    "            - jnp.sum(jnp.sin(x + 1) * jnp.sin(y)) * 2\n",
+    "        )"
    ]
   },
   {

docs/tutorials/linear/000_One_Sinkhorn.ipynb

@@ -19,19 +19,16 @@
    },
    "outputs": [],
    "source": [
-    "import functools\n",
     "import time\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "import ott\n",
-    "from ott import problems\n",
     "from ott.geometry import geometry, pointcloud\n",
     "from ott.solvers import linear\n",
-    "from ott.solvers.linear import acceleration, sinkhorn"
+    "from ott.solvers.linear import acceleration"
    ]
   },
   {
@@ -1087,7 +1084,7 @@
     "    plot_results(\n",
     "        out_scaling[i],\n",
     "        leg_scaling[i],\n",
-    "        title=rf\"Decay, $\\varepsilon$=\" + str(epsilon),\n",
+    "        title=r\"Decay, $\\varepsilon$=\" + str(epsilon),\n",
     "        xlabel=\"iterations\",\n",
     "        ylabel=\"error\",\n",
     "    )"
@@ -1276,7 +1273,7 @@
     "    plot_results(\n",
     "        out_mixed[i],\n",
     "        leg_mixed[i],\n",
-    "        title=rf\"Mixed strategy, $\\varepsilon$=\" + str(epsilon),\n",
+    "        title=r\"Mixed strategy, $\\varepsilon$=\" + str(epsilon),\n",
     "        xlabel=\"iterations\",\n",
     "        ylabel=\"error\",\n",
     "    )"

docs/tutorials/linear/100_OTT_&_POT.ipynb

@@ -23,21 +23,19 @@
    },
    "outputs": [],
    "source": [
-    "import timeit\n",
-    "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import numpy as np\n",
     "import ot\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
-    "\n",
-    "plt.rc(\"font\", size=20)\n",
     "import mpl_toolkits.axes_grid1\n",
     "\n",
     "from ott.geometry import pointcloud\n",
     "from ott.problems.linear import linear_problem\n",
-    "from ott.solvers.linear import sinkhorn"
+    "from ott.solvers.linear import sinkhorn\n",
+    "\n",
+    "plt.rc(\"font\", size=20)"
    ]
   },
   {

docs/tutorials/linear/200_sinkhorn_divergence_gradient_flow.ipynb

@@ -29,7 +29,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from typing import Any, Callable, Tuple\n",
+    "from typing import Any, Callable\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",

docs/tutorials/linear/600_mmsink.ipynb

@@ -15,8 +15,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import math\n",
-    "from typing import List, Optional, Tuple\n",
+    "from typing import Optional\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",

docs/tutorials/misc/000_tracking_progress.ipynb

@@ -29,15 +29,12 @@
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "import numpy as np\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
     "\n",
     "from ott import utils\n",
     "from ott.geometry import pointcloud\n",
     "from ott.problems.linear import linear_problem\n",
     "from ott.problems.quadratic import quadratic_problem\n",
-    "from ott.solvers import linear, quadratic\n",
+    "from ott.solvers import linear\n",
     "from ott.solvers.linear import sinkhorn, sinkhorn_lr\n",
     "from ott.solvers.quadratic import gromov_wasserstein"
    ]

docs/tutorials/misc/200_application_biology.ipynb

@@ -262,8 +262,8 @@
     "for day in DAYS:\n",
     "    tmp = cell_distribution_ipsc[day].copy()\n",
     "    tmp[tmp >= 1e-2] = alpha_bins[0]\n",
-    "    tmp[np.logical_and(1e-2 > tmp, tmp >= 5e-4)] = alpha_bins[1]\n",
-    "    tmp[5e-4 > tmp] = alpha_bins[2]\n",
+    "    tmp[np.logical_and(tmp < 1e-2, tmp >= 5e-4)] = alpha_bins[1]\n",
+    "    tmp[tmp < 5e-4] = alpha_bins[2]\n",
     "    binned_cell_distribution_ipsc[day] = tmp"
    ]
   },

docs/tutorials/neural/000_neural_dual.ipynb

@@ -25,8 +25,6 @@
    "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",
     "\n",

docs/tutorials/neural/200_Monge_Gap.ipynb

@@ -16,7 +16,7 @@
     "import dataclasses\n",
     "from collections.abc import Iterator, Mapping\n",
     "from types import MappingProxyType\n",
-    "from typing import Any, Dict, Literal, Optional, Tuple, Union\n",
+    "from typing import Any, Literal, Optional\n",
     "\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
@@ -31,7 +31,6 @@
     "from ott.geometry import costs, pointcloud\n",
     "from ott.neural.methods import monge_gap\n",
     "from ott.neural.networks import potentials\n",
-    "from ott.solvers.linear import acceleration\n",
     "from ott.tools import sinkhorn_divergence"
    ]
   },
@@ -587,7 +586,7 @@
     }
    ],
    "source": [
-    "plot_fit_map(rf\"with $\\ell_2$ Monge gap\", out_nn_l2, logs_l2)"
+    "plot_fit_map(r\"with $\\ell_2$ Monge gap\", out_nn_l2, logs_l2)"
    ]
   },
   {
@@ -659,7 +658,7 @@
     }
    ],
    "source": [
-    "plot_fit_map(rf\"with $\\ell_2^2$ Monge gap\", out_nn_l22, logs_l22)"
+    "plot_fit_map(r\"with $\\ell_2^2$ Monge gap\", out_nn_l22, logs_l22)"
    ]
   }
  ],

docs/tutorials/quadratic/000_gromov_wasserstein.ipynb

@@ -24,9 +24,8 @@
     "import numpy as np\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
-    "import mpl_toolkits.mplot3d.axes3d as p3\n",
     "from IPython import display\n",
-    "from matplotlib import animation, cm\n",
+    "from matplotlib import animation\n",
     "\n",
     "from ott.geometry import pointcloud\n",
     "from ott.problems.quadratic import quadratic_problem\n",

docs/utils.rst

@@ -11,3 +11,4 @@ function for :class:`~ott.solvers.linear.sinkhorn.Sinkhorn`.
 
     default_progress_fn
     tqdm_progress_fn
+    batched_vmap

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ott-jax"
-description = "Optimal Transport Tools in JAX."
+description = "Optimal Transport Tools in JAX"
 requires-python = ">=3.9"
 dynamic = ["version"]
 readme = {file = "README.md", content-type = "text/markdown"}
@@ -17,6 +17,7 @@ dependencies = [
     "jaxopt>=0.8",
     "lineax>=0.0.5",
     "numpy>=1.20.0",
+    "typing_extensions; python_version <= '3.9'",
 ]
 keywords = [
     "optimal transport",
@@ -107,7 +108,7 @@ multi_line_output = 3
 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", "ot", "torch", "torchvision", "pandas", "sklearn", "tslearn"]
+known_numeric = ["numpy", "scipy", "jax", "chex", "flax", "optax", "jaxopt", "ot", "torch", "torchvision", "pandas", "sklearn", "tslearn"]
 known_test = ["_pytest", "pytest"]
 known_plotting = ["IPython", "matplotlib", "mpl_toolkits", "seaborn"]
 
@@ -120,12 +121,6 @@ markers = [
     "cpu: Mark tests as CPU only.",
     "fast: Mark tests as fast.",
 ]
-filterwarnings = [
-    "ignore:\\n*.*scipy.sparse array",
-    "ignore:jax.random.KeyArray is deprecated:DeprecationWarning",
-    "ignore:.*jax.config:DeprecationWarning",
-    "ignore:jax.core.Shape is deprecated:DeprecationWarning:chex",
-]
 
 [tool.coverage.run]
 branch = true