diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 745717241..dcf260adf 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -38,7 +38,7 @@ repos: rev: 1.8.5 hooks: - id: nbqa-pyupgrade - args: [--py39-plus] + args: [--py311-plus] - id: nbqa-black - id: nbqa-isort - repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks diff --git a/docs/tutorials/basic_ot_between_datasets.ipynb b/docs/tutorials/basic_ot_between_datasets.ipynb index 5733a39d2..3ea12a089 100644 --- a/docs/tutorials/basic_ot_between_datasets.ipynb +++ b/docs/tutorials/basic_ot_between_datasets.ipynb @@ -31,7 +31,6 @@ "outputs": [], "source": [ "import jax\n", - "import jax.numpy as jnp\n", "\n", "import matplotlib.pyplot as plt\n", "\n", @@ -246,7 +245,7 @@ "metadata": {}, "outputs": [], "source": [ - "def reg_ot_cost(x: jnp.ndarray, y: jnp.ndarray) -> float:\n", + "def reg_ot_cost(x: jax.Array, y: jax.Array) -> float:\n", " geom = pointcloud.PointCloud(x, y)\n", " ot = solve_fn(geom)\n", " return ot.reg_ot_cost" diff --git a/docs/tutorials/geometry/000_point_cloud.ipynb b/docs/tutorials/geometry/000_point_cloud.ipynb index 2a6aac888..041c1aba9 100644 --- a/docs/tutorials/geometry/000_point_cloud.ipynb +++ b/docs/tutorials/geometry/000_point_cloud.ipynb @@ -223,8 +223,8 @@ "outputs": [], "source": [ "def optimize(\n", - " x: jnp.ndarray,\n", - " y: jnp.ndarray,\n", + " x: jax.Array,\n", + " y: jax.Array,\n", " num_iter: int = 300,\n", " dump_every: int = 5,\n", " learning_rate: float = 0.2,\n", diff --git a/docs/tutorials/geometry/100_grid.ipynb b/docs/tutorials/geometry/100_grid.ipynb index ed4161226..7742919c6 100644 --- a/docs/tutorials/geometry/100_grid.ipynb +++ b/docs/tutorials/geometry/100_grid.ipynb @@ -241,10 +241,10 @@ "class MyCost(costs.CostFn):\n", " \"\"\"An unusual cost function.\"\"\"\n", "\n", - " def norm(self, x: jnp.ndarray) -> jnp.ndarray:\n", + " def norm(self, x: jax.Array) -> jax.Array:\n", " return jnp.sum(x**3 + jnp.cos(x) ** 2, axis=-1)\n", "\n", - " def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray:\n", + " def __call__(self, x: jax.Array, y: jax.Array) -> jax.Array:\n", " return (\n", " self.norm(x)\n", " + self.norm(y)\n", diff --git a/docs/tutorials/linear/000_One_Sinkhorn.ipynb b/docs/tutorials/linear/000_One_Sinkhorn.ipynb index 3fd753b3f..3cbb97dd1 100644 --- a/docs/tutorials/linear/000_One_Sinkhorn.ipynb +++ b/docs/tutorials/linear/000_One_Sinkhorn.ipynb @@ -544,9 +544,7 @@ }, "outputs": [], "source": [ - "def my_sinkhorn(\n", - " geom: geometry.Geometry, a: jnp.ndarray, b: jnp.ndarray, **kwargs\n", - "):\n", + "def my_sinkhorn(geom: geometry.Geometry, a: jax.Array, b: jax.Array, **kwargs):\n", " return linear.solve(\n", " geom, a, b, inner_iterations=1, max_iterations=10_000, **kwargs\n", " )" diff --git a/docs/tutorials/linear/200_sinkhorn_divergence_gradient_flow.ipynb b/docs/tutorials/linear/200_sinkhorn_divergence_gradient_flow.ipynb index e7ff1405d..039598a02 100644 --- a/docs/tutorials/linear/200_sinkhorn_divergence_gradient_flow.ipynb +++ b/docs/tutorials/linear/200_sinkhorn_divergence_gradient_flow.ipynb @@ -30,7 +30,8 @@ "outputs": [], "source": [ "import functools\n", - "from typing import Any, Callable\n", + "from collections.abc import Callable\n", + "from typing import Any\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -116,9 +117,9 @@ "outputs": [], "source": [ "def gradient_flow(\n", - " x: jnp.ndarray,\n", - " y: jnp.ndarray,\n", - " divergence: Callable[[jnp.ndarray, jnp.ndarray, float], tuple[float, Any]],\n", + " x: jax.Array,\n", + " y: jax.Array,\n", + " divergence: Callable[[jax.Array, jax.Array, float], tuple[float, Any]],\n", " div_name: str = \"Obj\",\n", " num_iter: int = 500,\n", " lr: float = 0.2,\n", diff --git a/docs/tutorials/linear/400_Hessians.ipynb b/docs/tutorials/linear/400_Hessians.ipynb index ac38ff599..9616df10f 100644 --- a/docs/tutorials/linear/400_Hessians.ipynb +++ b/docs/tutorials/linear/400_Hessians.ipynb @@ -98,7 +98,7 @@ }, "outputs": [], "source": [ - "def loss(a: jnp.ndarray, x: jnp.ndarray, implicit: bool = True) -> float:\n", + "def loss(a: jax.Array, x: jax.Array, implicit: bool = True) -> float:\n", " div, _ = sinkhorn_divergence.sinkhorn_divergence(\n", " pointcloud.PointCloud,\n", " x,\n", diff --git a/docs/tutorials/linear/500_sparse_monge_displacements.ipynb b/docs/tutorials/linear/500_sparse_monge_displacements.ipynb index c37d768b6..66bfb40af 100644 --- a/docs/tutorials/linear/500_sparse_monge_displacements.ipynb +++ b/docs/tutorials/linear/500_sparse_monge_displacements.ipynb @@ -211,7 +211,7 @@ "solver = jax.jit(sinkhorn.Sinkhorn())\n", "\n", "\n", - "def entropic_map(x, y, cost_fn: costs.TICost) -> jnp.ndarray:\n", + "def entropic_map(x, y, cost_fn: costs.TICost) -> jax.Array:\n", " geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn)\n", " output = solver(linear_problem.LinearProblem(geom))\n", " dual_potentials = output.to_dual_potentials()\n", diff --git a/docs/tutorials/linear/600_mmsink.ipynb b/docs/tutorials/linear/600_mmsink.ipynb index 7fd05f4c1..959b4fec7 100644 --- a/docs/tutorials/linear/600_mmsink.ipynb +++ b/docs/tutorials/linear/600_mmsink.ipynb @@ -15,8 +15,6 @@ "metadata": {}, "outputs": [], "source": [ - "from typing import Optional\n", - "\n", "import jax\n", "import jax.numpy as jnp\n", "\n", @@ -93,7 +91,7 @@ "\n", "def plot_clouds(\n", " out: mmsinkhorn.MMSinkhornOutput,\n", - " top_k: Optional[int] = None,\n", + " top_k: int | None = None,\n", ") -> None:\n", " fig, ax = plt.subplots(figsize=(6, 5), tight_layout=True)\n", " plott = plot.PlotMM(fig=fig, ax=ax, cmap=cmap)\n", @@ -126,8 +124,8 @@ "source": [ "def display_animation(\n", " ots: list[mmsinkhorn.MMSinkhornOutput],\n", - " top_k: Optional[int] = None,\n", - " titles: Optional[list[str]] = None,\n", + " top_k: int | None = None,\n", + " titles: list[str] | None = None,\n", " frame_rate: int = 5,\n", ") -> None:\n", " fig = plt.figure(figsize=(6, 5))\n", @@ -202201,8 +202199,8 @@ "outputs": [], "source": [ "def objective(\n", - " x_s: list[jnp.ndarray],\n", - " a_s: Optional[list[jnp.ndarray]] = None,\n", + " x_s: list[jax.Array],\n", + " a_s: list[jax.Array] | None = None,\n", " epsilon: float = 1e-2,\n", ") -> tuple[float, mmsinkhorn.MMSinkhornOutput]:\n", " out = mmsinkhorn.MMSinkhorn()(x_s=x_s, a_s=a_s, epsilon=epsilon)\n", diff --git a/docs/tutorials/linear/700_progot.ipynb b/docs/tutorials/linear/700_progot.ipynb index 8d7245db7..4aa33fcc0 100644 --- a/docs/tutorials/linear/700_progot.ipynb +++ b/docs/tutorials/linear/700_progot.ipynb @@ -40,7 +40,7 @@ }, "outputs": [], "source": [ - "from typing import Any, Optional\n", + "from typing import Any\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -262,10 +262,10 @@ ], "source": [ "def entropic_map(\n", - " x: jnp.ndarray,\n", - " y: jnp.ndarray,\n", + " x: jax.Array,\n", + " y: jax.Array,\n", " cost_fn: costs.TICost,\n", - " epsilon: Optional[float] = None,\n", + " epsilon: float | None = None,\n", ") -> potentials.DualPotentials:\n", " geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn, epsilon=epsilon)\n", " output = linear.solve(geom)\n", @@ -313,7 +313,7 @@ "outputs": [], "source": [ "def run_progot(\n", - " x: jnp.ndarray, y: jnp.ndarray, cost_fn: costs.TICost, **kwargs: Any\n", + " x: jax.Array, y: jax.Array, cost_fn: costs.TICost, **kwargs: Any\n", ") -> progot.ProgOTOutput:\n", " geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn)\n", " prob = linear_problem.LinearProblem(geom)\n", diff --git a/docs/tutorials/linear/800_Unbalanced_OT.ipynb b/docs/tutorials/linear/800_Unbalanced_OT.ipynb index bf766744e..f19947c22 100644 --- a/docs/tutorials/linear/800_Unbalanced_OT.ipynb +++ b/docs/tutorials/linear/800_Unbalanced_OT.ipynb @@ -55,8 +55,8 @@ "outputs": [], "source": [ "def generate_data(\n", - " rng: jax.Array, *, means: jnp.ndarray, cov: jnp.ndarray, n_samples: int\n", - ") -> jnp.ndarray:\n", + " rng: jax.Array, *, means: jax.Array, cov: jax.Array, n_samples: int\n", + ") -> jax.Array:\n", " gmm = gaussian_mixture.GaussianMixture.from_mean_cov_component_weights(\n", " mean=means,\n", " cov=cov,\n", diff --git a/docs/tutorials/linear/900_CIFAR_benchmark.ipynb b/docs/tutorials/linear/900_CIFAR_benchmark.ipynb index 73dc4806a..6133c7bfd 100644 --- a/docs/tutorials/linear/900_CIFAR_benchmark.ipynb +++ b/docs/tutorials/linear/900_CIFAR_benchmark.ipynb @@ -18,7 +18,6 @@ "# Misc.\n", "import functools\n", "import time\n", - "from typing import Optional\n", "\n", "# Jax and sharding utils\n", "import jax\n", @@ -177,7 +176,7 @@ "from jax.sharding import PartitionSpec as P\n", "\n", "\n", - "def shard_samples(x: jnp.ndarray) -> tuple[jnp.ndarray, NamedSharding]:\n", + "def shard_samples(x: jax.Array) -> tuple[jax.Array, NamedSharding]:\n", " num_devices = jax.device_count()\n", " devices = mesh_utils.create_device_mesh((num_devices,))\n", " mesh = Mesh(devices, axis_names=(\"batch\",))\n", @@ -230,10 +229,10 @@ " jax.jit, static_argnames=(\"batch_size\",), in_shardings=in_shardings\n", ")\n", "def run(\n", - " x: jnp.ndarray,\n", - " y: jnp.ndarray,\n", - " epsilon: Optional[float] = None,\n", - " batch_size: Optional[int] = None,\n", + " x: jax.Array,\n", + " y: jax.Array,\n", + " epsilon: float | None = None,\n", + " batch_size: int | None = None,\n", ") -> sinkhorn.SinkhornOutput:\n", " geom = pointcloud.PointCloud(\n", " x,\n", diff --git a/docs/tutorials/misc/300_otcp.ipynb b/docs/tutorials/misc/300_otcp.ipynb index d834aaa75..f9129a6cd 100644 --- a/docs/tutorials/misc/300_otcp.ipynb +++ b/docs/tutorials/misc/300_otcp.ipynb @@ -46,7 +46,7 @@ "metadata": {}, "outputs": [], "source": [ - "from typing import Any, Optional\n", + "from typing import Any\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -308,9 +308,9 @@ "outputs": [], "source": [ "def plot(\n", - " x: jnp.ndarray,\n", - " y: jnp.ndarray,\n", - " label: Optional[str] = None,\n", + " x: jax.Array,\n", + " y: jax.Array,\n", + " label: str | None = None,\n", " ax=None,\n", " **kwargs: Any,\n", ") -> None:\n", @@ -877,7 +877,7 @@ } ], "source": [ - "def model_fn(x: jnp.ndarray) -> jnp.ndarray:\n", + "def model_fn(x: jax.Array) -> jax.Array:\n", " return jax.pure_callback(skl_model.predict, jnp.empty_like(x), x)\n", "\n", "\n", diff --git a/docs/tutorials/neural/000_neural_dual.ipynb b/docs/tutorials/neural/000_neural_dual.ipynb index 404caad92..82afc8c65 100644 --- a/docs/tutorials/neural/000_neural_dual.ipynb +++ b/docs/tutorials/neural/000_neural_dual.ipynb @@ -370,9 +370,7 @@ "outputs": [], "source": [ "@jax.jit\n", - "def sinkhorn_loss(\n", - " x: jnp.ndarray, y: jnp.ndarray, epsilon: float = 0.1\n", - ") -> float:\n", + "def sinkhorn_loss(x: jax.Array, y: jax.Array, epsilon: float = 0.1) -> float:\n", " \"\"\"Computes transport between (x, a) and (y, b) via Sinkhorn algorithm.\"\"\"\n", " a = jnp.ones(len(x)) / len(x)\n", " b = jnp.ones(len(y)) / len(y)\n", diff --git a/docs/tutorials/neural/200_Monge_Gap.ipynb b/docs/tutorials/neural/200_Monge_Gap.ipynb index 02eee209a..4be5cdfd7 100644 --- a/docs/tutorials/neural/200_Monge_Gap.ipynb +++ b/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, Literal, Optional\n", + "from typing import Any, Literal\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -84,13 +84,13 @@ "\n", " name: Literal[\"moon\", \"s_curve\"]\n", " theta_rotation: float = 0.0\n", - " mean: Optional[jnp.ndarray] = None\n", + " mean: jax.Array | None = None\n", " noise: float = 0.01\n", " scale: float = 1.0\n", " batch_size: int = 1024\n", - " rng: Optional[jax.Array] = None\n", + " rng: jax.Array | None = None\n", "\n", - " def __iter__(self) -> Iterator[jnp.ndarray]:\n", + " def __iter__(self) -> Iterator[jax.Array]:\n", " \"\"\"Random sample generator from Gaussian mixture.\n", "\n", " Returns:\n", @@ -98,7 +98,7 @@ " \"\"\"\n", " return self._create_sample_generators()\n", "\n", - " def _create_sample_generators(self) -> Iterator[jnp.ndarray]:\n", + " def _create_sample_generators(self) -> Iterator[jax.Array]:\n", " rng = jax.random.key(0) if self.rng is None else self.rng\n", "\n", " # define rotation matrix tp rotate samples\n", @@ -141,7 +141,7 @@ " target_kwargs: Mapping[str, Any] = MappingProxyType({}),\n", " train_batch_size: int = 256,\n", " valid_batch_size: int = 256,\n", - " rng: Optional[jax.Array] = None,\n", + " rng: jax.Array | None = None,\n", ") -> tuple[datasets.Dataset, datasets.Dataset, int]:\n", " \"\"\"Samplers from ``SklearnDistribution``.\"\"\"\n", " rng = jax.random.key(0) if rng is None else rng\n", @@ -189,10 +189,10 @@ "source": [ "def plot_samples(\n", " batch: dict[str, Any],\n", - " num_points: Optional[int] = None,\n", - " title: Optional[str] = None,\n", + " num_points: int | None = None,\n", + " title: str | None = None,\n", " figsize: tuple[int, int] = (8, 6),\n", - " rng: Optional[jax.Array] = None,\n", + " rng: jax.Array | None = None,\n", "):\n", " \"\"\"Plot samples from the source and target measures.\n", "\n", @@ -471,7 +471,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 5000/5000 [00:54<00:00, 92.53it/s, fitting_loss: 0.1896, regularizer: NA ,total: 0.1896] \n" + "100%|██████████| 5000/5000 [00:54<00:00, 92.53it/s, fitting_loss: 0.1896, regularizer: NA ,total: 0.1896] \n" ] } ], @@ -542,7 +542,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 5000/5000 [00:59<00:00, 83.46it/s, fitting_loss: 0.2151, regularizer: 0.7066 ,total: 1.6283] \n" + "100%|██████████| 5000/5000 [00:59<00:00, 83.46it/s, fitting_loss: 0.2151, regularizer: 0.7066 ,total: 1.6283] \n" ] } ], @@ -614,7 +614,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "100%|\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 5000/5000 [01:23<00:00, 60.16it/s, fitting_loss: 0.1948, regularizer: 0.6860 ,total: 0.2634]\n" + "100%|██████████| 5000/5000 [01:23<00:00, 60.16it/s, fitting_loss: 0.1948, regularizer: 0.6860 ,total: 0.2634]\n" ] } ], diff --git a/docs/tutorials/neural/300_ENOT.ipynb b/docs/tutorials/neural/300_ENOT.ipynb index d653c661e..55bbb1986 100644 --- a/docs/tutorials/neural/300_ENOT.ipynb +++ b/docs/tutorials/neural/300_ENOT.ipynb @@ -28,7 +28,7 @@ "import dataclasses\n", "from collections.abc import Iterator, Mapping\n", "from types import MappingProxyType\n", - "from typing import Any, Literal, Optional\n", + "from typing import Any, Literal\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -301,9 +301,7 @@ "outputs": [], "source": [ "@jax.jit\n", - "def sinkhorn_loss(\n", - " x: jnp.ndarray, y: jnp.ndarray, epsilon: float = 0.001\n", - ") -> float:\n", + "def sinkhorn_loss(x: jax.Array, y: jax.Array, epsilon: float = 0.001) -> float:\n", " \"\"\"Computes transport between (x, a) and (y, b) via Sinkhorn algorithm.\"\"\"\n", " a = jnp.ones(len(x)) / len(x)\n", " b = jnp.ones(len(y)) / len(y)\n", @@ -561,16 +559,16 @@ "class SklearnDistribution:\n", " name: Literal[\"moon\", \"s_curve\"]\n", " theta_rotation: float = 0.0\n", - " mean: Optional[jnp.ndarray] = None\n", + " mean: jax.Array | None = None\n", " noise: float = 0.01\n", " scale: float = 1.0\n", " batch_size: int = 1024\n", - " rng: Optional[jax.Array] = None\n", + " rng: jax.Array | None = None\n", "\n", - " def __iter__(self) -> Iterator[jnp.ndarray]:\n", + " def __iter__(self) -> Iterator[jax.Array]:\n", " return self._create_sample_generators()\n", "\n", - " def _create_sample_generators(self) -> Iterator[jnp.ndarray]:\n", + " def _create_sample_generators(self) -> Iterator[jax.Array]:\n", " rng = jax.random.key(0) if self.rng is None else self.rng\n", " rotation = jnp.array(\n", " [\n", @@ -611,7 +609,7 @@ " target_kwargs: Mapping[str, Any] = MappingProxyType({}),\n", " train_batch_size: int = 256,\n", " valid_batch_size: int = 256,\n", - " rng: Optional[jax.Array] = None,\n", + " rng: jax.Array | None = None,\n", "):\n", " rng = jax.random.key(0) if rng is None else rng\n", " rng1, rng2, rng3, rng4 = jax.random.split(rng, 4)\n", diff --git a/docs/tutorials/neural/400_MetaOT.ipynb b/docs/tutorials/neural/400_MetaOT.ipynb index 93d67e9db..d87a519c8 100644 --- a/docs/tutorials/neural/400_MetaOT.ipynb +++ b/docs/tutorials/neural/400_MetaOT.ipynb @@ -408,7 +408,7 @@ " num_hidden_layers: int = 3\n", "\n", " @nn.compact\n", - " def __call__(self, z: jnp.ndarray) -> jnp.ndarray:\n", + " def __call__(self, z: jax.Array) -> jax.Array:\n", " for _ in range(self.num_hidden_layers):\n", " z = nn.relu(nn.Dense(self.num_hidden_units)(z))\n", "\n", diff --git a/docs/tutorials/neural/500_otfm.ipynb b/docs/tutorials/neural/500_otfm.ipynb index 224dcef3a..a1a3eac5c 100644 --- a/docs/tutorials/neural/500_otfm.ipynb +++ b/docs/tutorials/neural/500_otfm.ipynb @@ -19,8 +19,7 @@ "outputs": [], "source": [ "# Basic imports\n", - "from collections.abc import Iterable\n", - "from typing import Callable, Dict, Literal, Optional, Tuple\n", + "from collections.abc import Callable, Iterable\n", "\n", "from tqdm.auto import trange\n", "\n", @@ -91,7 +90,7 @@ "\n", "\n", "@jax.jit\n", - "def phi(x: jnp.ndarray) -> jnp.ndarray:\n", + "def phi(x: jax.Array) -> jax.Array:\n", " \"\"\"Real-valued convex potential function.\"\"\"\n", " return (\n", " 2 * jnp.sum(jnp.abs(A @ x) ** 1.7)\n", @@ -191,7 +190,7 @@ "def gen_points(\n", " rng: jax.Array,\n", " batch_size: tuple[int, ...],\n", - " dtype: Optional[jnp.dtype] = None,\n", + " dtype: jnp.dtype | None = None,\n", ") -> jax.Array:\n", " batch_size, *_ = batch_size\n", " x = 3 * jnp.array(gen_torus_points(rng, size=batch_size // 2), dtype=dtype)\n", @@ -19971,8 +19970,8 @@ "def unpaired_dl(\n", " rng: jax.Array,\n", " batch_size: int,\n", - " potential: Callable[[jnp.ndarray], jnp.ndarray],\n", - ") -> Iterable[tuple[jnp.ndarray, jnp.ndarray]]:\n", + " potential: Callable[[jax.Array], jax.Array],\n", + ") -> Iterable[tuple[jax.Array, jax.Array]]:\n", " while True:\n", " rng, rng_x0, rng_x1 = jr.split(rng, 3)\n", " x0 = gen_points(rng_x0, (batch_size,))\n", diff --git a/docs/tutorials/neural/600_otfm_2.ipynb b/docs/tutorials/neural/600_otfm_2.ipynb index 85442d15d..0d45e3bee 100644 --- a/docs/tutorials/neural/600_otfm_2.ipynb +++ b/docs/tutorials/neural/600_otfm_2.ipynb @@ -39,7 +39,6 @@ "source": [ "import pickle\n", "from collections.abc import Iterable\n", - "from typing import Callable, Dict, Literal, Optional, Tuple\n", "\n", "# For loading images\n", "from PIL import Image, ImageOps\n", diff --git a/docs/tutorials/theory/000_bb_flow.ipynb b/docs/tutorials/theory/000_bb_flow.ipynb index 695d88a90..cebb9d85a 100644 --- a/docs/tutorials/theory/000_bb_flow.ipynb +++ b/docs/tutorials/theory/000_bb_flow.ipynb @@ -26,7 +26,6 @@ "outputs": [], "source": [ "import functools\n", - "from typing import Callable, Dict, Optional\n", "\n", "import jax\n", "import jax.numpy as jnp\n", @@ -80,7 +79,7 @@ "assert jnp.linalg.det(B) > 0.0\n", "\n", "\n", - "def phi(x: jnp.ndarray) -> jnp.ndarray:\n", + "def phi(x: jax.Array) -> jax.Array:\n", " \"\"\"Real-valued convex potential function.\"\"\"\n", " return (\n", " jnp.sum(jnp.log(1.0 + jnp.exp(B @ x)))\n", diff --git a/pyproject.toml b/pyproject.toml index 9b34c613b..d7228ceb0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -185,7 +185,7 @@ exclude = [ "dist" ] line-length = 80 -target-version = "py38" # TODO(michalk8): bump to py311, see #XXX +target-version = "py311" [tool.ruff.lint] ignore = [ @@ -201,6 +201,10 @@ ignore = [ "D107", # Missing docstring in magic method "D105", + # deprecated by ruff; `isinstance(x, X | Y)` is slower than the tuple form + "UP038", + # `zip()` without `strict=` -> several call sites rely on truncation + "B905", ] select = [ "D", # flake8-docstrings diff --git a/src/ott/data/_loaders.py b/src/ott/data/_loaders.py index a99e1a544..269c2a964 100644 --- a/src/ott/data/_loaders.py +++ b/src/ott/data/_loaders.py @@ -22,7 +22,7 @@ def get_cifar10( root: str = "data/cifar10", batch_size: int = 1000, use_flip: bool = True -) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array, jax.Array]: transform = transforms.Compose([ transforms.ToTensor(), @@ -53,7 +53,7 @@ def get_cifar10( target_batches = [] label_batches = [] - def psd_gaussian_blur(x: jnp.ndarray) -> jnp.ndarray: + def psd_gaussian_blur(x: jax.Array) -> jax.Array: assert x.ndim == 4, x.shape n, c, h, w = x.shape assert h == w diff --git a/src/ott/datasets.py b/src/ott/datasets.py index e19ac7182..c01aec6d8 100644 --- a/src/ott/datasets.py +++ b/src/ott/datasets.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Iterator, Literal, NamedTuple, Optional, Tuple +from collections.abc import Iterator +from typing import Literal, NamedTuple import jax import jax.numpy as jnp @@ -39,8 +40,8 @@ class Dataset(NamedTuple): target_iter: loader for the target measure """ - source_iter: Iterator[jnp.ndarray] - target_iter: Iterator[jnp.ndarray] + source_iter: Iterator[jax.Array] + target_iter: Iterator[jax.Array] class ConditionalDataset(NamedTuple): @@ -54,10 +55,10 @@ class ConditionalDataset(NamedTuple): label_iter: loader for integer condition labels, ``[batch]`` """ - source_iter: Iterator[jnp.ndarray] - target_iter: Iterator[jnp.ndarray] - condition_iter: Iterator[jnp.ndarray] - label_iter: Iterator[jnp.ndarray] + source_iter: Iterator[jax.Array] + target_iter: Iterator[jax.Array] + condition_iter: Iterator[jax.Array] + label_iter: Iterator[jax.Array] @dataclasses.dataclass @@ -135,8 +136,8 @@ def create_gaussian_mixture_samplers( name_target: Name_t, train_batch_size: int = 2048, valid_batch_size: int = 2048, - rng: Optional[jax.Array] = None, -) -> Tuple[Dataset, Dataset, int]: + rng: jax.Array | None = None, +) -> tuple[Dataset, Dataset, int]: """Gaussian samplers. Args: @@ -190,13 +191,13 @@ class ConditionalGaussianMixture: num_conditions: int batch_size: int dim: int - offsets: jnp.ndarray + offsets: jax.Array rng: jax.Array - def __iter__(self) -> Iterator[Tuple[jnp.ndarray, ...]]: + def __iter__(self) -> Iterator[tuple[jax.Array, ...]]: return self._generate() - def _generate(self) -> Iterator[Tuple[jnp.ndarray, ...]]: + def _generate(self) -> Iterator[tuple[jax.Array, ...]]: rng = self.rng per_cond = self.batch_size // self.num_conditions while True: @@ -225,8 +226,8 @@ def create_conditional_gaussian_mixture_samplers( dim: int = 2, train_batch_size: int = 90, valid_batch_size: int = 90, - rng: Optional[jax.Array] = None, -) -> Tuple[ConditionalDataset, ConditionalDataset, int, int, int]: + rng: jax.Array | None = None, +) -> tuple[ConditionalDataset, ConditionalDataset, int, int, int]: """Create conditional Gaussian samplers for testing. Each condition defines a different translation of the source distribution. @@ -270,7 +271,7 @@ def _next_batch(): cache["batch"] = next(gen) return cache - def _iter(idx: int) -> Iterator[jnp.ndarray]: + def _iter(idx: int) -> Iterator[jax.Array]: while True: c = _next_batch() val = c["batch"][idx] diff --git a/src/ott/experimental/mmsinkhorn.py b/src/ott/experimental/mmsinkhorn.py index 9c037f5f5..883faa1fb 100644 --- a/src/ott/experimental/mmsinkhorn.py +++ b/src/ott/experimental/mmsinkhorn.py @@ -11,7 +11,7 @@ # 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 typing import Any, NamedTuple, Optional, Tuple, Union +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -26,13 +26,13 @@ class MMSinkhornState(NamedTuple): - potentials: Tuple[jnp.ndarray, ...] - errors: jnp.ndarray + potentials: tuple[jax.Array, ...] + errors: jax.Array def solution_error( self, - cost_t: jnp.ndarray, - a_s: Tuple[jnp.ndarray, ...], + cost_t: jax.Array, + a_s: tuple[jax.Array, ...], epsilon: float, norm_error: float = 1.0 ) -> float: @@ -80,16 +80,16 @@ class MMSinkhornOutput(NamedTuple): inner_iterations: Number of iterations that were run between two computations of errors. """ - potentials: Tuple[jnp.ndarray, ...] - errors: jnp.ndarray - x_s: Optional[Tuple[jnp.ndarray, ...]] = None - a_s: Optional[Tuple[jnp.ndarray, ...]] = None - cost_fns: Optional[Union[costs.CostFn, Tuple[costs.CostFn, ...]]] = None - epsilon: Optional[float] = None - ent_reg_cost: Optional[jnp.ndarray] = None - threshold: Optional[jnp.ndarray] = None - converged: Optional[bool] = None - inner_iterations: Optional[int] = None + potentials: tuple[jax.Array, ...] + errors: jax.Array + x_s: tuple[jax.Array, ...] | None = None + a_s: tuple[jax.Array, ...] | None = None + cost_fns: costs.CostFn | tuple[costs.CostFn, ...] | None = None + epsilon: float | None = None + ent_reg_cost: jax.Array | None = None + threshold: jax.Array | None = None + converged: bool | None = None + inner_iterations: int | None = None def set(self, **kwargs: Any) -> "MMSinkhornOutput": """Return a copy of self, with potential overwrites.""" @@ -101,23 +101,23 @@ def n_iters(self) -> int: # noqa: D102 return jnp.sum(self.errors != -1) * self.inner_iterations @property - def cost_t(self) -> jnp.ndarray: + def cost_t(self) -> jax.Array: """Cost tensor.""" return cost_tensor(self.x_s, self.cost_fns) @property - def tensor(self) -> jnp.ndarray: + def tensor(self) -> jax.Array: """Transport tensor.""" return jnp.exp( -remove_tensor_sum(self.cost_t, self.potentials) / self.epsilon ) @property - def marginals(self) -> Tuple[jnp.ndarray, ...]: + def marginals(self) -> tuple[jax.Array, ...]: """:math:`k` marginal probability weight vectors.""" return tensor_marginals(self.tensor) - def marginal(self, k: int) -> jnp.ndarray: + def marginal(self, k: int) -> jax.Array: """Return the marginal probability weight vector at slice :math:`k`.""" return tensor_marginal(self.tensor, k) @@ -127,7 +127,7 @@ def transport_mass(self) -> float: return jnp.sum(self.tensor) @property - def shape(self) -> Tuple[int, ...]: + def shape(self) -> tuple[int, ...]: """Shape of the transport :attr:`tensor`.""" return tuple(x.shape[0] for x in self.x_s) @@ -138,9 +138,9 @@ def n_marginals(self) -> int: def cost_tensor( - x_s: Tuple[jnp.ndarray, ...], cost_fns: Union[costs.CostFn, - Tuple[costs.CostFn, ...]] -) -> jnp.ndarray: + x_s: tuple[jax.Array, ...], + cost_fns: costs.CostFn | tuple[costs.CostFn, ...] +) -> jax.Array: r"""Create a cost tensor from a tuple of :math:`k` :math:`d`-dim point clouds. Args: @@ -171,9 +171,7 @@ def c_fn_pair(i: int, j: int) -> costs.CostFn: return cost_t -def remove_tensor_sum( - c: jnp.ndarray, u: Tuple[jnp.ndarray, ...] -) -> jnp.ndarray: +def remove_tensor_sum(c: jax.Array, u: tuple[jax.Array, ...]) -> jax.Array: r"""Remove the tensor sum of :math:`k` vectors to tensor of :math:`k` dims. Args: @@ -189,11 +187,11 @@ def remove_tensor_sum( return c -def tensor_marginals(coupling: jnp.ndarray) -> Tuple[jnp.ndarray, ...]: +def tensor_marginals(coupling: jax.Array) -> tuple[jax.Array, ...]: return tuple(tensor_marginal(coupling, ix) for ix in range(coupling.ndim)) -def tensor_marginal(coupling: jnp.ndarray, slice_index: int) -> jnp.ndarray: +def tensor_marginal(coupling: jax.Array, slice_index: int) -> jax.Array: k = coupling.ndim axis = list(range(slice_index)) + list(range(slice_index + 1, k)) return coupling.sum(axis=axis) @@ -251,10 +249,10 @@ def __init__( def __call__( self, - x_s: Tuple[jnp.ndarray, ...], - a_s: Optional[Tuple[jnp.ndarray, ...]] = None, - cost_fns: Optional[Union[costs.CostFn, Tuple[costs.CostFn, ...]]] = None, - epsilon: Optional[float] = None + x_s: tuple[jax.Array, ...], + a_s: tuple[jax.Array, ...] | None = None, + cost_fns: costs.CostFn | tuple[costs.CostFn, ...] | None = None, + epsilon: float | None = None ) -> MMSinkhornOutput: r"""Solve multimarginal OT for :math:`k` :math:`d`-dim point clouds. @@ -288,7 +286,7 @@ def __call__( n_s = [x.shape[0] for x in x_s] if cost_fns is None: cost_fns = costs.SqEuclidean() - elif isinstance(cost_fns, Tuple): + elif isinstance(cost_fns, tuple): assert len(cost_fns) == (len(n_s) * (len(n_s) - 1)) // 2 # Default to uniform probability weights for each point cloud. @@ -310,7 +308,7 @@ def __call__( out = run(const, self, state) return out.set(x_s=x_s, a_s=a_s, cost_fns=cost_fns, epsilon=epsilon) - def init_state(self, n_s: Tuple[int, ...]) -> MMSinkhornState: + def init_state(self, n_s: tuple[int, ...]) -> MMSinkhornState: """Return the initial state of the loop.""" errors = -jnp.ones((self.outer_iterations, 1)) potentials = tuple(jnp.zeros(n) for n in n_s) @@ -351,25 +349,25 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 def run( - const: Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...], float], - solver: MMSinkhorn, state: MMSinkhornState + const: tuple[jax.Array, tuple[jax.Array, ...], float], solver: MMSinkhorn, + state: MMSinkhornState ) -> MMSinkhornOutput: def cond_fn( - iteration: int, const: Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...], float], + iteration: int, const: tuple[jax.Array, tuple[jax.Array, ...], float], state: MMSinkhornState ) -> bool: del const return solver._continue(state, iteration) def body_fn( - iteration: int, const: Tuple[jnp.ndarray, Tuple[jnp.ndarray, ...], float], + iteration: int, const: tuple[jax.Array, tuple[jax.Array, ...], float], state: MMSinkhornState, compute_error: bool ) -> MMSinkhornState: cost_t, a_s, epsilon = const k = len(a_s) - def one_slice(potentials: Tuple[jnp.ndarray, ...], l: int, a: jnp.ndarray): + def one_slice(potentials: tuple[jax.Array, ...], l: int, a: jax.Array): pot = potentials[l] axis = list(range(l)) + list(range(l + 1, k)) app_lse = mu.softmin( @@ -431,6 +429,6 @@ def one_slice(potentials: Tuple[jnp.ndarray, ...], l: int, a: jnp.ndarray): def coupling_tensor( - potentials: Tuple[jnp.ndarray], cost_t: jnp.ndarray, epsilon: float -) -> jnp.ndarray: + potentials: tuple[jax.Array], cost_t: jax.Array, epsilon: float +) -> jax.Array: return jnp.exp(-remove_tensor_sum(cost_t, potentials) / epsilon) diff --git a/src/ott/geometry/costs.py b/src/ott/geometry/costs.py index 3a2e8d69e..7b6eb3c6e 100644 --- a/src/ott/geometry/costs.py +++ b/src/ott/geometry/costs.py @@ -14,7 +14,8 @@ import abc import functools import math -from typing import Any, Callable, Dict, Optional, Tuple +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -33,7 +34,7 @@ ] # TODO(michalk8): norm check -Func = Callable[[jnp.ndarray], float] +Func = Callable[[jax.Array], float] @jtu.register_pytree_node_class @@ -41,7 +42,7 @@ class CostFn(abc.ABC): """Base class for all costs.""" @abc.abstractmethod - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute cost between :math:`x` and :math:`y`. Args: @@ -52,8 +53,8 @@ def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: The cost. """ - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + def barycenter(self, weights: jax.Array, + xs: jax.Array) -> tuple[jax.Array, Any]: """Barycentric operator. Args: @@ -68,7 +69,7 @@ def barycenter(self, weights: jnp.ndarray, raise NotImplementedError("Barycenter is not implemented.") @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: + def _padder(cls, dim: int) -> jax.Array: """Create a padding vector of adequate dimension, well-suited to a cost. Args: @@ -79,7 +80,7 @@ def _padder(cls, dim: int) -> jnp.ndarray: """ return jnp.zeros((1, dim)) - def all_pairs(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + def all_pairs(self, x: jax.Array, y: jax.Array) -> jax.Array: """Compute matrix of all pairwise costs, including the :attr:`norms `. Args: @@ -92,8 +93,8 @@ def all_pairs(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: return jax.vmap(lambda x_: jax.vmap(lambda y_: self(x_, y_))(y))(x) def twist_operator( - self, vec: jnp.ndarray, dual_vec: jnp.ndarray, variable: bool - ) -> jnp.ndarray: + self, vec: jax.Array, dual_vec: jax.Array, variable: bool + ) -> jax.Array: r"""Twist inverse operator of the cost function. Given a cost function :math:`c`, the twist operator returns @@ -138,7 +139,7 @@ class TICost(CostFn): """ @abc.abstractmethod - def h(self, z: jnp.ndarray) -> float: + def h(self, z: jax.Array) -> float: """TI function acting on difference of :math:`x-y` to output cost. Args: @@ -148,20 +149,20 @@ def h(self, z: jnp.ndarray) -> float: The cost. """ - def h_legendre(self, z: jnp.ndarray) -> float: + def h_legendre(self, z: jax.Array) -> float: """Legendre transform of :func:`h` when it is convex.""" raise NotImplementedError("Legendre transform of `h` is not implemented.") - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute cost as evaluation of :func:`h` on :math:`x-y`.""" return self.h(x - y) def h_transform( self, f: Func, - solver: Optional[Callable[[Func, jnp.ndarray, jnp.ndarray, Any], - jnp.ndarray]] = None, - ) -> Callable[[jnp.ndarray, Optional[jnp.ndarray], Any], float]: + solver: Callable[[Func, jax.Array, jax.Array, Any], jax.Array] + | None = None, + ) -> Callable[[jax.Array, jax.Array | None, Any], float]: r"""Compute the h-transform of a concave function. Return a callable :math:`f_h` defined as: @@ -191,9 +192,7 @@ def h_transform( solver = ott_math.lbfgs def f_h( - x: jnp.ndarray, - x_init: Optional[jnp.ndarray] = None, - **kwargs: Any + x: jax.Array, x_init: jax.Array | None = None, **kwargs: Any ) -> float: """h-transform of a concave function. @@ -207,7 +206,7 @@ def f_h( The :math:`h`-transform of :math:`f`, :math:`f_h(x)`. """ - def fun(z: jnp.ndarray) -> float: + def fun(z: jax.Array) -> float: return self.h(z) - f(x - z) x_init = x if x_init is None else x_init @@ -218,14 +217,14 @@ def fun(z: jnp.ndarray) -> float: return f_h def twist_operator( - self, vec: jnp.ndarray, dual_vec: jnp.ndarray, variable: bool - ) -> jnp.ndarray: + self, vec: jax.Array, dual_vec: jax.Array, variable: bool + ) -> jax.Array: # Note: when `h` is pair, i.e. h(z) = h(-z), the expressions below coincide if variable: return vec + jax.grad(self.h_legendre)(-dual_vec) return vec - jax.grad(self.h_legendre)(dual_vec) - def transport_map(self, g: Func) -> Callable[[jnp.ndarray, Any], jnp.ndarray]: + def transport_map(self, g: Func) -> Callable[[jax.Array, Any], jax.Array]: r"""Get an optimal transport map for a concave function :math:`g`. Uses Proposition 1 from :cite:`klein:24` to define an OT map @@ -240,7 +239,7 @@ def transport_map(self, g: Func) -> Callable[[jnp.ndarray, Any], jnp.ndarray]: The transport map with a signature ``(x, **kwargs)``. """ - def transport(x: jnp.ndarray, **kwargs: Any) -> jnp.ndarray: + def transport(x: jax.Array, **kwargs: Any) -> jax.Array: """Transport points from source to target. Args: @@ -259,8 +258,8 @@ def transport(x: jnp.ndarray, **kwargs: Any) -> jnp.ndarray: return transport - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + def barycenter(self, weights: jax.Array, + xs: jax.Array) -> tuple[jax.Array, Any]: """Output barycenter of vectors.""" return jnp.average(xs, weights=weights, axis=0), None @@ -281,10 +280,10 @@ def __init__(self, p: float): self.p = p self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf - def h(self, z: jnp.ndarray) -> float: # noqa: D102 + def h(self, z: jax.Array) -> float: # noqa: D102 return 0.5 * mu.norm(z, self.p) ** 2 - def h_legendre(self, z: jnp.ndarray) -> float: + def h_legendre(self, z: jax.Array) -> float: """Legendre transform of :func:`h`. For details on the derivation, see e.g., :cite:`boyd:04`, p. 93/94. @@ -317,10 +316,10 @@ def __init__(self, p: float): self.p = p self.q = 1.0 / (1.0 - (1.0 / p)) if p > 1.0 else jnp.inf - def h(self, z: jnp.ndarray) -> float: # noqa: D102 + def h(self, z: jax.Array) -> float: # noqa: D102 return mu.norm(z, self.p) ** self.p / self.p - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + def h_legendre(self, z: jax.Array) -> float: # noqa: D102 # not defined for `p=1` return mu.norm(z, self.q) ** self.q / self.q @@ -348,7 +347,7 @@ def __init__(self, p: float): super().__init__() self.p = p - def h(self, z: jnp.ndarray) -> float: # noqa: D102 + def h(self, z: jax.Array) -> float: # noqa: D102 return mu.norm(z, ord=2) ** self.p def tree_flatten(self): # noqa: D102 @@ -370,20 +369,20 @@ class NegDotProduct(CostFn): c(x,y) = - \langle x, y\rangle """ - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array, y: jax.Array) -> float: # noqa: D102 return -jnp.vdot(x, y) - def twist_operator(self, vec, dual_vec, variable) -> jnp.ndarray: + def twist_operator(self, vec, dual_vec, variable) -> jax.Array: """Twist operator for negative dot-product cost.""" del vec, variable return -dual_vec - def norm(self, x: jnp.ndarray) -> jnp.ndarray: + def norm(self, x: jax.Array) -> jax.Array: """Compute squared Euclidean norm for vector. Only used for rescaling.""" return jnp.sum(x ** 2, axis=-1) - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + def barycenter(self, weights: jax.Array, + xs: jax.Array) -> tuple[jax.Array, Any]: """Output usual barycenter of vectors.""" return jnp.average(xs, weights=weights, axis=0), None @@ -416,10 +415,10 @@ def __init__( rho=rho, ) - def h(self, z: jnp.ndarray) -> float: # noqa: D102 + def h(self, z: jax.Array) -> float: # noqa: D102 return self._h(z) - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + def h_legendre(self, z: jax.Array) -> float: # noqa: D102 """Legendre transform of :func:`h`. This function uses :class:`~jax.custom_vjp` to apply Danskin's theorem @@ -433,15 +432,15 @@ def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 """ @jax.custom_vjp - def fn(z: jnp.ndarray) -> float: + def fn(z: jax.Array) -> float: out, _ = fwd(z) return out - def fwd(z: jnp.ndarray) -> Tuple[float, jnp.ndarray]: + def fwd(z: jax.Array) -> tuple[float, jax.Array]: q = self.regularizer.prox(z) return jnp.dot(q, z) - self.h(q), q - def bwd(q: jnp.ndarray, g: jnp.ndarray) -> Tuple[jnp.ndarray]: + def bwd(q: jax.Array, g: jax.Array) -> tuple[jax.Array]: return jnp.dot(g, q), fn.defvjp(fwd, bwd) @@ -450,7 +449,7 @@ def bwd(q: jnp.ndarray, g: jnp.ndarray) -> Tuple[jnp.ndarray]: def h_transform( self, f: Func, - ) -> Callable[[jnp.ndarray, Optional[jnp.ndarray], Any], float]: + ) -> Callable[[jax.Array, jax.Array | None, Any], float]: r"""Compute the h-transform of a concave function. Return a callable :math:`f_h` defined as: @@ -480,9 +479,7 @@ def h_transform( """ def f_h( - x: jnp.ndarray, - x_init: Optional[jnp.ndarray] = None, - **kwargs: Any + x: jax.Array, x_init: jax.Array | None = None, **kwargs: Any ) -> float: """h-transform of a concave function. @@ -539,7 +536,7 @@ class Euclidean(CostFn): because the function is not strictly convex (it is linear on rays). """ - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute sq. Euclidean distance using a custom jvp implementation. Here we use a custom jvp implementation for the norm that does not yield @@ -556,23 +553,23 @@ class SqEuclidean(TICost): Implemented as a translation invariant cost, :math:`h(z) = \|z\|^2`. """ - def norm(self, x: jnp.ndarray) -> jnp.ndarray: + def norm(self, x: jax.Array) -> jax.Array: """Compute squared Euclidean norm for vector.""" return jnp.sum(x ** 2, axis=-1) - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute minus twice the dot-product between vectors.""" cross_term = -2.0 * jnp.vdot(x, y) return self.norm(x) + self.norm(y) + cross_term - def h(self, z: jnp.ndarray) -> float: # noqa: D102 + def h(self, z: jax.Array) -> float: # noqa: D102 return jnp.sum(z ** 2) - def h_legendre(self, z: jnp.ndarray) -> float: # noqa: D102 + def h_legendre(self, z: jax.Array) -> float: # noqa: D102 return 0.25 * jnp.sum(z ** 2) - def barycenter(self, weights: jnp.ndarray, - xs: jnp.ndarray) -> Tuple[jnp.ndarray, Any]: + def barycenter(self, weights: jax.Array, + xs: jax.Array) -> tuple[jax.Array, Any]: """Output barycenter of vectors when using squared-Euclidean distance.""" return jnp.average(xs, weights=weights, axis=0), None @@ -589,7 +586,7 @@ def __init__(self, ridge: float = 1e-8): super().__init__() self._ridge = ridge - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Cosine distance between vectors, denominator regularized with ridge.""" x_norm = jnp.linalg.norm(x, axis=-1) y_norm = jnp.linalg.norm(y, axis=-1) @@ -597,7 +594,7 @@ def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: return 1.0 - cosine_similarity @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: + def _padder(cls, dim: int) -> jax.Array: return jnp.ones((1, dim)) @@ -625,7 +622,7 @@ def __init__(self, n: int, ridge: float = 1e-8): self.n = n self._ridge = ridge - def __call__(self, x: jnp.ndarray, y: jnp.ndarray): # noqa: D102 + def __call__(self, x: jax.Array, y: jax.Array): # noqa: D102 x_norm = jnp.linalg.norm(x, axis=-1) y_norm = jnp.linalg.norm(y, axis=-1) cosine_similarity = jnp.vdot(x, y) / (x_norm * y_norm + self._ridge) @@ -677,19 +674,19 @@ class Bures(CostFn): behavior of inner calls to :func:`~ott.math.matrix_square_root.sqrtm`. """ - def __init__(self, dimension: int, sqrtm_kw: Optional[Dict[str, Any]] = None): + def __init__(self, dimension: int, sqrtm_kw: dict[str, Any] | None = None): super().__init__() self._dimension = dimension self._sqrtm_kw = {} if sqrtm_kw is None else sqrtm_kw - def norm(self, x: jnp.ndarray) -> jnp.ndarray: + def norm(self, x: jax.Array) -> jax.Array: """Compute norm of Gaussian, sq. 2-norm of mean + trace of covariance.""" mean, cov = x_to_means_and_covs(x, self._dimension) norm = jnp.sum(mean ** 2, axis=-1) norm += jnp.trace(cov, axis1=-2, axis2=-1) return norm - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute - 2 x Bures dot-product.""" mean_x, cov_x = x_to_means_and_covs(x, self._dimension) mean_y, cov_y = x_to_means_and_covs(y, self._dimension) @@ -706,12 +703,12 @@ def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: def covariance_fixpoint_iter( self, - covs: jnp.ndarray, - weights: jnp.ndarray, + covs: jax.Array, + weights: jax.Array, tolerance: float = 1e-4, - sqrtm_kw: Optional[Dict[str, Any]] = None, + sqrtm_kw: dict[str, Any] | None = None, **kwargs: Any - ) -> jnp.ndarray: + ) -> jax.Array: """Iterate fix-point updates to compute barycenter of Gaussians. Args: @@ -736,21 +733,21 @@ def covariance_fixpoint_iter( @functools.partial(jax.vmap, in_axes=[None, 0, 0]) def scale_covariances( - cov_sqrt: jnp.ndarray, cov: jnp.ndarray, weight: jnp.ndarray - ) -> jnp.ndarray: + cov_sqrt: jax.Array, cov: jax.Array, weight: jax.Array + ) -> jax.Array: """Rescale covariance in barycenter step.""" return weight * matrix_square_root.sqrtm_only((cov_sqrt @ cov) @ cov_sqrt, **sqrtm_kw) - def cond_fn(iteration: int, constants: Tuple[Any, ...], state) -> bool: + def cond_fn(iteration: int, constants: tuple[Any, ...], state) -> bool: del constants _, diffs = state return diffs[iteration // inner_iterations] > tolerance def body_fn( - iteration: int, constants: Tuple[Any, ...], - state: Tuple[jnp.ndarray, float], compute_error: bool - ) -> Tuple[jnp.ndarray, float]: + iteration: int, constants: tuple[Any, ...], + state: tuple[jax.Array, float], compute_error: bool + ) -> tuple[jax.Array, float]: del constants, compute_error cov, diffs = state cov_sqrt, cov_inv_sqrt, _ = matrix_square_root.sqrtm(cov, **sqrtm_kw) @@ -762,7 +759,7 @@ def body_fn( diffs = diffs.at[iteration // inner_iterations].set(diff) return next_cov, diffs - def init_state() -> Tuple[jnp.ndarray, float]: + def init_state() -> tuple[jax.Array, float]: cov_init = jnp.eye(self._dimension) diffs = -jnp.ones(math.ceil(max_iterations / inner_iterations)) return cov_init, diffs @@ -780,12 +777,12 @@ def init_state() -> Tuple[jnp.ndarray, float]: def barycenter( self, - weights: jnp.ndarray, - xs: jnp.ndarray, + weights: jax.Array, + xs: jax.Array, tolerance: float = 1e-4, - sqrtm_kw: Optional[Dict[str, Any]] = None, + sqrtm_kw: dict[str, Any] | None = None, **kwargs: Any - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array]: """Compute the Bures barycenter of weighted Gaussian distributions. Implements the fixed point approach proposed in :cite:`alvarez-esteban:16` @@ -831,7 +828,7 @@ def barycenter( return mean_and_cov_to_x(mu_bary, cov_bary, self._dimension), diffs @classmethod - def _padder(cls, dim: int) -> jnp.ndarray: + def _padder(cls, dim: int) -> jax.Array: dimension = int((-1 + math.sqrt(1 + 4 * dim)) / 2) padding = mean_and_cov_to_x( jnp.zeros((dimension,)), jnp.eye(dimension), dimension @@ -874,7 +871,7 @@ def __init__( self._gamma = gamma self._sqrtm_kw = kwargs - def norm(self, x: jnp.ndarray) -> jnp.ndarray: + def norm(self, x: jax.Array) -> jax.Array: """Compute norm of Gaussian for unbalanced Bures. Args: @@ -887,7 +884,7 @@ def norm(self, x: jnp.ndarray) -> jnp.ndarray: """ return self._gamma * x[..., 0] - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Compute dot-product for unbalanced Bures. Args: @@ -975,25 +972,24 @@ class SoftDTW(CostFn): def __init__( self, gamma: float, - ground_cost: Optional[CostFn] = None, + ground_cost: CostFn | None = None, debiased: bool = False ): self.gamma = gamma self.ground_cost = SqEuclidean() if ground_cost is None else ground_cost self.debiased = debiased - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array, y: jax.Array) -> float: # noqa: D102 c_xy = self._soft_dtw(x, y) if self.debiased: return c_xy - 0.5 * (self._soft_dtw(x, x) + self._soft_dtw(y, y)) return c_xy - def _soft_dtw(self, t1: jnp.ndarray, t2: jnp.ndarray) -> float: + def _soft_dtw(self, t1: jax.Array, t2: jax.Array) -> float: def body( - carry: Tuple[jnp.ndarray, jnp.ndarray], - current_antidiagonal: jnp.ndarray - ) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], jnp.ndarray]: + carry: tuple[jax.Array, jax.Array], current_antidiagonal: jax.Array + ) -> tuple[tuple[jax.Array, jax.Array], jax.Array]: # modified from: https://github.com/khdlr/softdtw_jax two_ago, one_ago = carry @@ -1040,8 +1036,8 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 return cls(*children, **aux_data) -def x_to_means_and_covs(x: jnp.ndarray, - dimension: int) -> Tuple[jnp.ndarray, jnp.ndarray]: +def x_to_means_and_covs(x: jax.Array, + dimension: int) -> tuple[jax.Array, jax.Array]: """Extract means and covariance matrices of Gaussians from raveled vector. Args: @@ -1061,8 +1057,8 @@ def x_to_means_and_covs(x: jnp.ndarray, def mean_and_cov_to_x( - mean: jnp.ndarray, covariance: jnp.ndarray, dimension: int -) -> jnp.ndarray: + mean: jax.Array, covariance: jax.Array, dimension: int +) -> jax.Array: """Ravel a Gaussian's mean and covariance matrix to d(1 + d) vector.""" return jnp.concatenate( (mean, jnp.reshape(covariance, (dimension * dimension))) diff --git a/src/ott/geometry/distrib_costs.py b/src/ott/geometry/distrib_costs.py index d1adb6055..7251a6b34 100644 --- a/src/ott/geometry/distrib_costs.py +++ b/src/ott/geometry/distrib_costs.py @@ -11,8 +11,9 @@ # 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 typing import Callable, Optional +from collections.abc import Callable +import jax import jax.numpy as jnp import jax.tree_util as jtu @@ -43,7 +44,7 @@ def __init__( self, solve_fn: Callable[[linear_problem.LinearProblem], univariate.UnivariateOutput], - ground_cost: Optional[costs.TICost] = None, + ground_cost: costs.TICost | None = None, ): super().__init__() self.ground_cost = ( @@ -51,7 +52,7 @@ def __init__( ) self._solve_fn = solve_fn - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: """Wasserstein distance between :math:`x` and :math:`y` seen as a 1D dist. Args: diff --git a/src/ott/geometry/epsilon_scheduler.py b/src/ott/geometry/epsilon_scheduler.py index bebb3d051..d612f1090 100644 --- a/src/ott/geometry/epsilon_scheduler.py +++ b/src/ott/geometry/epsilon_scheduler.py @@ -11,7 +11,6 @@ # 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 typing import Optional import jax.numpy as jnp import jax.tree_util as jtu @@ -44,7 +43,7 @@ def __init__(self, target: jnp.array, init: float = 1.0, decay: float = 1.0): self.init = init self.decay = decay - def __call__(self, it: Optional[int]) -> jnp.array: + def __call__(self, it: int | None) -> jnp.array: """Intermediate regularizer value at a given iteration number. Args: diff --git a/src/ott/geometry/geodesic.py b/src/ott/geometry/geodesic.py index 2b11f635c..d57a7bea5 100644 --- a/src/ott/geometry/geodesic.py +++ b/src/ott/geometry/geodesic.py @@ -11,13 +11,15 @@ # 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 typing import Any, Dict, Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Any import jax import jax.experimental.sparse as jesp import jax.numpy as jnp import jax.tree_util as jtu import numpy as np +from jax.typing import DTypeLike from scipy.special import ive from ott import utils @@ -26,7 +28,7 @@ __all__ = ["Geodesic"] -Array_g = Union[jnp.ndarray, jesp.BCOO] +Array_g = jax.Array | jesp.BCOO @jtu.register_pytree_node_class @@ -53,8 +55,8 @@ class Geodesic(geometry.Geometry): def __init__( self, scaled_laplacian: Array_g, - eigval: jnp.ndarray, - chebyshev_coeffs: jnp.ndarray, + eigval: jax.Array, + chebyshev_coeffs: jax.Array, t: float = 1e-3, **kwargs: Any ): @@ -68,12 +70,12 @@ def __init__( def from_graph( cls, G: Array_g, - t: Optional[float] = 1e-3, - eigval: Optional[jnp.ndarray] = None, + t: float | None = 1e-3, + eigval: jax.Array | None = None, order: int = 100, directed: bool = False, normalize: bool = False, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, **kwargs: Any ) -> "Geodesic": r"""Construct a Geodesic geometry from an adjacency matrix. @@ -135,10 +137,10 @@ def from_graph( def apply_kernel( self, - vec: jnp.ndarray, - eps: Optional[float] = None, + vec: jax.Array, + eps: float | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: r"""Apply :attr:`kernel_matrix` on a positive vector. Args: @@ -154,7 +156,7 @@ def apply_kernel( ) @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + def kernel_matrix(self) -> jax.Array: # noqa: D102 n, _ = self.shape kernel = self.apply_kernel(jnp.eye(n)) return jax.lax.cond( @@ -163,12 +165,12 @@ def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 ) @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + def cost_matrix(self) -> jax.Array: # noqa: D102 # Calculate the cost matrix using the formula (5) from the main reference return -4.0 * self.t * mu.safe_log(self.kernel_matrix) @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self.scaled_laplacian.shape @property @@ -179,32 +181,30 @@ def is_symmetric(self) -> bool: # noqa: D102 def dtype(self) -> jnp.dtype: # noqa: D102 return self.scaled_laplacian.dtype - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: + def transport_from_potentials(self, f: jax.Array, g: jax.Array) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") def apply_transport_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, + f: jax.Array, + g: jax.Array, + vec: jax.Array, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") def marginal_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [ self.scaled_laplacian, self.eigval, @@ -214,19 +214,17 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "Geodesic": return cls(*children, **aux_data) -def normalize_laplacian(laplacian: Array_g, degree: jnp.ndarray) -> Array_g: +def normalize_laplacian(laplacian: Array_g, degree: jax.Array) -> Array_g: inv_sqrt_deg = jnp.where(degree > 0.0, 1.0 / jnp.sqrt(degree), 0.0) return inv_sqrt_deg[:, None] * laplacian * inv_sqrt_deg[None, :] -def compute_dense_laplacian( - G: jnp.ndarray, normalize: bool = False -) -> jnp.ndarray: +def compute_dense_laplacian(G: jax.Array, normalize: bool = False) -> jax.Array: degree = jnp.sum(G, axis=1) laplacian = jnp.diag(degree) - G if normalize: @@ -253,7 +251,7 @@ def compute_sparse_laplacian( def compute_largest_eigenvalue( - laplacian_matrix: jnp.ndarray, + laplacian_matrix: jax.Array, rng: jax.Array, ) -> float: # Compute the largest eigenvalue of the Laplacian matrix. @@ -274,8 +272,8 @@ def compute_largest_eigenvalue( def expm_multiply( - L: Array_g, X: jnp.ndarray, coeff: jnp.ndarray, eigval: float -) -> jnp.ndarray: + L: Array_g, X: jax.Array, coeff: jax.Array, eigval: float +) -> jax.Array: def body(carry, c): T0, T1, Y = carry @@ -294,8 +292,8 @@ def body(carry, c): def compute_chebychev_coeff_all( - eigval: float, tau: float, K: int, dtype: np.dtype -) -> jnp.ndarray: + eigval: float, tau: float, K: int, dtype: DTypeLike +) -> jax.Array: """Jax wrapper to compute the K+1 Chebychev coefficients.""" result_shape_dtype = jax.ShapeDtypeStruct( shape=(K + 1,), diff --git a/src/ott/geometry/geometry.py b/src/ott/geometry/geometry.py index e51ba95b8..269df7b28 100644 --- a/src/ott/geometry/geometry.py +++ b/src/ott/geometry/geometry.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Tuple, Union +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Literal if TYPE_CHECKING: from ott.geometry import low_rank @@ -75,12 +76,11 @@ class Geometry: def __init__( self, - cost_matrix: Optional[jnp.ndarray] = None, - kernel_matrix: Optional[jnp.ndarray] = None, - epsilon: Optional[Union[float, eps_scheduler.Epsilon]] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_cost", "median", - "std"]] = 1.0, + cost_matrix: jax.Array | None = None, + kernel_matrix: jax.Array | None = None, + epsilon: float | eps_scheduler.Epsilon | None = None, + relative_epsilon: Literal["mean", "std"] | None = None, + scale_cost: float | Literal["mean", "max_cost", "median", "std"] = 1.0, ): self._cost_matrix = cost_matrix self._kernel_matrix = kernel_matrix @@ -89,11 +89,11 @@ def __init__( self._scale_cost = scale_cost @property - def cost_rank(self) -> Optional[int]: + def cost_rank(self) -> int | None: """Output rank of cost matrix, if any was provided.""" @property - def cost_matrix(self) -> jnp.ndarray: + def cost_matrix(self) -> jax.Array: """Cost matrix, recomputed from kernel if only kernel was specified.""" if self._cost_matrix is None: # If no epsilon was passed on to the geometry, then assume it is one by @@ -135,7 +135,7 @@ def std_cost_matrix(self) -> float: return jnp.sqrt(jax.nn.relu(tmp)) @property - def kernel_matrix(self) -> jnp.ndarray: + def kernel_matrix(self) -> jax.Array: """Kernel matrix. Either provided by user or recomputed from :attr:`cost_matrix`. @@ -176,7 +176,7 @@ def epsilon(self) -> float: return self.epsilon_scheduler.target @property - def shape(self) -> Tuple[int, int]: + def shape(self) -> tuple[int, int]: """Shape of the geometry.""" mat = ( self._kernel_matrix if self._cost_matrix is None else self._cost_matrix @@ -217,7 +217,7 @@ def is_square(self) -> bool: return (n == m) @property - def inv_scale_cost(self) -> jnp.ndarray: + def inv_scale_cost(self) -> jax.Array: """Compute and return inverse of scaling factor for cost matrix.""" if self._scale_cost == "max_cost": return 1.0 / jnp.max(self._cost_matrix) @@ -230,12 +230,12 @@ def inv_scale_cost(self) -> jnp.ndarray: raise ValueError(f"Scaling {self._scale_cost} not implemented.") @property - def diag_cost(self) -> jnp.ndarray: + def diag_cost(self) -> jax.Array: """Diagonal of the cost matrix.""" assert self.is_square, "Cost matrix must be square to compute diagonal." return jnp.diag(self.cost_matrix) - def set_scale_cost(self, scale_cost: Union[float, str]) -> "Geometry": + def set_scale_cost(self, scale_cost: float | str) -> "Geometry": """Modify how to rescale of the :attr:`cost_matrix`.""" # case when `geom` doesn't have `scale_cost` or doesn't need to be modified # `False` retains the original scale @@ -259,12 +259,12 @@ def copy_epsilon(self, other: "Geometry") -> "Geometry": def apply_lse_kernel( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, eps: float, - vec: jnp.ndarray = None, + vec: jax.Array = None, axis: int = 0 - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array]: r"""Apply :attr:`kernel_matrix` in log domain. This function applies the ground geometry's kernel in log domain, using @@ -281,10 +281,10 @@ def apply_lse_kernel( f and g in iterations 1 & 2 respectively. Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + f: jax.Array [num_a,] , potential of size num_rows of cost_matrix + g: jax.Array [num_b,] , potential of size num_cols of cost_matrix eps: float, regularization strength - vec: jnp.ndarray [num_a or num_b,] , when not None, this has the effect of + vec: jax.Array [num_a or num_b,] , when not None, this has the effect of doing log-Kernel computations with an addition elementwise multiplication of exp(g / eps) by a vector. This is carried out by adding weights to the log-sum-exp function, and needs to handle signs @@ -292,7 +292,7 @@ def apply_lse_kernel( axis: summing over axis 0 when doing (2), or over axis 1 when doing (1) Returns: - A jnp.ndarray corresponding to output above, depending on axis. + A jax.Array corresponding to output above, depending on axis. """ w_res, w_sgn = self._softmax(f, g, eps, vec, axis) remove = f if axis == 1 else g @@ -300,20 +300,20 @@ def apply_lse_kernel( def apply_kernel( self, - vec: jnp.ndarray, - eps: Optional[float] = None, + vec: jax.Array, + eps: float | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Apply :attr:`kernel_matrix` on positive scaling vector. Args: - vec: jnp.ndarray [num_a or num_b] , scaling of size num_rows or + vec: jax.Array [num_a or num_b] , scaling of size num_rows or num_cols of kernel_matrix eps: passed for consistency, not used yet. axis: standard kernel product if axis is 1, transpose if 0. Returns: - a jnp.ndarray corresponding to output above, depending on axis. + a jax.Array corresponding to output above, depending on axis. """ if eps is None: kernel = self.kernel_matrix @@ -325,10 +325,10 @@ def apply_kernel( def marginal_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Output marginal of transportation matrix from potentials. This applies first lse kernel in the standard way, removes the @@ -337,8 +337,8 @@ def marginal_from_potentials( by potentials. Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + f: jax.Array [num_a,] , potential of size num_rows of cost_matrix + g: jax.Array [num_b,] , potential of size num_cols of cost_matrix axis: axis along which to integrate, returns marginal on other axis. Returns: @@ -350,23 +350,19 @@ def marginal_from_potentials( def marginal_from_scalings( self, - u: jnp.ndarray, - v: jnp.ndarray, + u: jax.Array, + v: jax.Array, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Output marginal of transportation matrix from scalings.""" u, v = (v, u) if axis == 0 else (u, v) return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: + def transport_from_potentials(self, f: jax.Array, g: jax.Array) -> jax.Array: """Output transport matrix from potentials.""" return jnp.exp(self._center(f, g) / self.epsilon) - def transport_from_scalings( - self, u: jnp.ndarray, v: jnp.ndarray - ) -> jnp.ndarray: + def transport_from_scalings(self, u: jax.Array, v: jax.Array) -> jax.Array: """Output transport matrix from pair of scalings.""" return self.kernel_matrix * u[:, jnp.newaxis] * v[jnp.newaxis, :] @@ -375,17 +371,17 @@ def transport_from_scalings( def update_potential( self, - f: jnp.ndarray, - g: jnp.ndarray, - log_marginal: jnp.ndarray, - iteration: Optional[int] = None, + f: jax.Array, + g: jax.Array, + log_marginal: jax.Array, + iteration: int | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Carry out one Sinkhorn update for potentials, i.e. in log space. Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix + f: jax.Array [num_a,] , potential of size num_rows of cost_matrix + g: jax.Array [num_b,] , potential of size num_cols of cost_matrix log_marginal: targeted marginal iteration: used to compute epsilon from schedule, if provided. axis: axis along which the update should be carried out. @@ -399,15 +395,15 @@ def update_potential( def update_scaling( self, - scaling: jnp.ndarray, - marginal: jnp.ndarray, - iteration: Optional[int] = None, + scaling: jax.Array, + marginal: jax.Array, + iteration: int | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Carry out one Sinkhorn update for scalings, using kernel directly. Args: - scaling: jnp.ndarray of num_a or num_b positive values. + scaling: jax.Array of num_a or num_b positive values. marginal: targeted marginal iteration: used to compute epsilon from schedule, if provided. axis: axis along which the update should be carried out. @@ -420,13 +416,13 @@ def update_scaling( return marginal / jnp.where(app_kernel > 0, app_kernel, 1.0) # Helper functions - def _center(self, f: jnp.ndarray, g: jnp.ndarray) -> jnp.ndarray: + def _center(self, f: jax.Array, g: jax.Array) -> jax.Array: return f[:, jnp.newaxis] + g[jnp.newaxis, :] - self.cost_matrix def _softmax( - self, f: jnp.ndarray, g: jnp.ndarray, eps: float, - vec: Optional[jnp.ndarray], axis: int - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + self, f: jax.Array, g: jax.Array, eps: float, vec: jax.Array | None, + axis: int + ) -> tuple[jax.Array, jax.Array]: """Apply softmax row or column wise, weighted by vec.""" if vec is not None: if axis == 0: @@ -443,8 +439,8 @@ def _softmax( @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) def _apply_transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray, vec: jnp.ndarray, axis: int - ) -> jnp.ndarray: + self, f: jax.Array, g: jax.Array, vec: jax.Array, axis: int + ) -> jax.Array: """Apply lse_kernel to arbitrary vector while keeping track of signs.""" lse_res, lse_sgn = self.apply_lse_kernel( f, g, self.epsilon, vec=vec, axis=axis @@ -455,11 +451,11 @@ def _apply_transport_from_potentials( # wrapper to allow default option for axis. def apply_transport_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, + f: jax.Array, + g: jax.Array, + vec: jax.Array, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: """Apply transport matrix computed from potentials to a (batched) vec. This approach does not instantiate the transport matrix itself, but uses @@ -470,9 +466,9 @@ def apply_transport_from_potentials( (b=..., return_sign=True) optional parameters of logsumexp. Args: - f: jnp.ndarray [num_a,] , potential of size num_rows of cost_matrix - g: jnp.ndarray [num_b,] , potential of size num_cols of cost_matrix - vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied + f: jax.Array [num_a,] , potential of size num_rows of cost_matrix + g: jax.Array [num_b,] , potential of size num_cols of cost_matrix + vec: jax.Array [batch, num_a or num_b], vector that will be multiplied by transport matrix corresponding to potentials f, g, and geom. axis: axis to differentiate left (0) or right (1) multiply. @@ -487,7 +483,7 @@ def apply_transport_from_potentials( @functools.partial(jax.vmap, in_axes=[None, None, None, 0, None]) def _apply_transport_from_scalings( - self, u: jnp.ndarray, v: jnp.ndarray, vec: jnp.ndarray, axis: int + self, u: jax.Array, v: jax.Array, vec: jax.Array, axis: int ): u, v = (u, v * vec) if axis == 1 else (v, u * vec) return u * self.apply_kernel(v, eps=self.epsilon, axis=axis) @@ -495,20 +491,20 @@ def _apply_transport_from_scalings( # wrapper to allow default option for axis def apply_transport_from_scalings( self, - u: jnp.ndarray, - v: jnp.ndarray, - vec: jnp.ndarray, + u: jax.Array, + v: jax.Array, + vec: jax.Array, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: """Apply transport matrix computed from scalings to a (batched) vec. This approach does not instantiate the transport matrix itself, but relies instead on the apply_kernel function. Args: - u: jnp.ndarray [num_a,] , scaling of size num_rows of cost_matrix - v: jnp.ndarray [num_b,] , scaling of size num_cols of cost_matrix - vec: jnp.ndarray [batch, num_a or num_b], vector that will be multiplied + u: jax.Array [num_a,] , scaling of size num_rows of cost_matrix + v: jax.Array [num_b,] , scaling of size num_cols of cost_matrix + vec: jax.Array [batch, num_a or num_b], vector that will be multiplied by transport matrix corresponding to scalings u, v, and geom. axis: axis to differentiate left (0) or right (1) multiply. @@ -521,7 +517,7 @@ def apply_transport_from_scalings( )[0, :] return self._apply_transport_from_scalings(u, v, vec, axis) - def potential_from_scaling(self, scaling: jnp.ndarray) -> jnp.ndarray: + def potential_from_scaling(self, scaling: jax.Array) -> jax.Array: """Compute dual potential vector from scaling vector. Args: @@ -532,7 +528,7 @@ def potential_from_scaling(self, scaling: jnp.ndarray) -> jnp.ndarray: """ return self.epsilon * jnp.log(scaling) - def scaling_from_potential(self, potential: jnp.ndarray) -> jnp.ndarray: + def scaling_from_potential(self, potential: jax.Array) -> jax.Array: """Compute scaling vector from dual potential. Args: @@ -546,7 +542,7 @@ def scaling_from_potential(self, potential: jnp.ndarray) -> jnp.ndarray: finite, jnp.exp(jnp.where(finite, potential / self.epsilon, 0.0)), 0.0 ) - def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + def apply_square_cost(self, arr: jax.Array, axis: int = 0) -> jax.Array: """Apply elementwise-square of cost matrix to array (vector or matrix). This function applies the ground geometry's cost matrix, to perform either @@ -567,11 +563,11 @@ def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: def apply_cost( self, - arr: jnp.ndarray, + arr: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + fn: Callable[[jax.Array], jax.Array] | None = None, is_linear: bool = False, - ) -> jnp.ndarray: + ) -> jax.Array: """Apply :attr:`cost_matrix` to array (vector or matrix). This function applies the ground geometry's cost matrix, to perform either @@ -580,7 +576,7 @@ def apply_cost( where C is [num_a, num_b] Args: - arr: jnp.ndarray [num_a or num_b, p], vector that will be multiplied by + arr: jax.Array [num_a or num_b, p], vector that will be multiplied by the cost matrix. axis: standard cost matrix if axis=1, transpose if 0 fn: function to apply to cost matrix element-wise before the dot product @@ -598,22 +594,22 @@ def apply_cost( def _apply_cost_to_vec( self, - vec: jnp.ndarray, + vec: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + fn: Callable[[jax.Array], jax.Array] | None = None, is_linear: bool = False, - ) -> jnp.ndarray: + ) -> jax.Array: """Apply ``[num_a, num_b]`` fn(cost) (or transpose) to vector. Args: - vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector + vec: jax.Array [num_a,] ([num_b,] if axis=1) vector axis: axis on which the reduction is done. fn: function optionally applied to cost matrix element-wise, before the doc product is_linear: Whether ``fn`` is linear. Returns: - A jnp.ndarray corresponding to cost x vector + A jax.Array corresponding to cost x vector """ del is_linear matrix = self.cost_matrix.T if axis == 0 else self.cost_matrix @@ -627,7 +623,7 @@ def prepare_divergences( *args: Any, static_b: bool = False, **kwargs: Any - ) -> Tuple["Geometry", ...]: + ) -> tuple["Geometry", ...]: """Instantiate 2 (or 3) geometries to compute a Sinkhorn divergence.""" size = 2 if static_b else 3 nones = [None, None, None] @@ -643,7 +639,7 @@ def to_LRCGeometry( self, rank: int = 0, tol: float = 1e-2, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, scale: float = 1.0 ) -> "low_rank.LRCGeometry": r"""Factorize the cost matrix using either SVD (full) or :cite:`indyk:19`. @@ -737,8 +733,8 @@ def to_LRCGeometry( def subset( self, - row_ixs: Optional[jnp.ndarray] = None, - col_ixs: Optional[jnp.ndarray] = None + row_ixs: jax.Array | None = None, + col_ixs: jax.Array | None = None ) -> "Geometry": """Subset rows or columns of a geometry. diff --git a/src/ott/geometry/graph.py b/src/ott/geometry/graph.py index a0cbbc0d5..fd43f9a37 100644 --- a/src/ott/geometry/graph.py +++ b/src/ott/geometry/graph.py @@ -11,7 +11,8 @@ # 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 typing import Any, Dict, Literal, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -49,7 +50,7 @@ class Graph(geometry.Geometry): def __init__( self, - laplacian: jnp.ndarray, + laplacian: jax.Array, t: float = 1e-3, n_steps: int = 100, numerical_scheme: Literal["backward_euler", @@ -67,8 +68,8 @@ def __init__( @classmethod def from_graph( cls, - G: jnp.ndarray, - t: Optional[float] = 1e-3, + G: jax.Array, + t: float | None = 1e-3, directed: bool = False, normalize: bool = False, **kwargs: Any @@ -114,10 +115,10 @@ def from_graph( def apply_kernel( self, - vec: jnp.ndarray, - eps: Optional[float] = None, + vec: jax.Array, + eps: float | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: r"""Apply :attr:`kernel_matrix` on a positive vector. Args: @@ -130,8 +131,8 @@ def apply_kernel( """ def conf_fn( - iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], - old_new: Tuple[jnp.ndarray, jnp.ndarray] + iteration: int, consts: tuple[jax.Array, jax.Array | None], + old_new: tuple[jax.Array, jax.Array] ) -> bool: del iteration, consts @@ -144,9 +145,9 @@ def conf_fn( return (jnp.nanmax(f) - jnp.nanmin(f)) > self.tol def body_fn( - iteration: int, consts: Tuple[jnp.ndarray, Optional[jnp.ndarray]], - old_new: Tuple[jnp.ndarray, jnp.ndarray], compute_errors: bool - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + iteration: int, consts: tuple[jax.Array, jax.Array | None], + old_new: tuple[jax.Array, jax.Array], compute_errors: bool + ) -> tuple[jax.Array, jax.Array]: del iteration, compute_errors L, scaled_lap = consts @@ -187,7 +188,7 @@ def body_fn( )[1] @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + def kernel_matrix(self) -> jax.Array: # noqa: D102 n, _ = self.shape kernel = self.apply_kernel(jnp.eye(n)) # Symmetrize the kernel if needed. Numerical imprecision @@ -198,7 +199,7 @@ def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 ) @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + def cost_matrix(self) -> jax.Array: # noqa: D102 return -self.t * mu.safe_log(self.kernel_matrix) @property @@ -213,17 +214,17 @@ def _scale(self) -> float: ) @property - def _scaled_laplacian(self) -> jnp.ndarray: + def _scaled_laplacian(self) -> jax.Array: """Laplacian scaled by a constant, depending on the numerical scheme.""" return self._scale * self.laplacian @property - def _M(self) -> jnp.ndarray: + def _M(self) -> jax.Array: n, _ = self.shape return self._scaled_laplacian + jnp.eye(n) @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self.laplacian.shape @property @@ -234,32 +235,30 @@ def is_symmetric(self) -> bool: # noqa: D102 def dtype(self) -> jnp.dtype: # noqa: D102 return self.laplacian.dtype - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: + def transport_from_potentials(self, f: jax.Array, g: jax.Array) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") def apply_transport_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, - vec: jnp.ndarray, + f: jax.Array, + g: jax.Array, + vec: jax.Array, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") def marginal_from_potentials( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [self.laplacian, self.t], { "n_steps": self.n_steps, "numerical_scheme": self.numerical_scheme, @@ -268,6 +267,6 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "Graph": return cls(*children, **aux_data) diff --git a/src/ott/geometry/grid.py b/src/ott/geometry/grid.py index 714175cb8..8d3f97cec 100644 --- a/src/ott/geometry/grid.py +++ b/src/ott/geometry/grid.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import itertools -from typing import Any, Callable, List, NoReturn, Optional, Sequence, Tuple +from collections.abc import Callable, Sequence +from typing import Any, NoReturn import jax import jax.numpy as jnp @@ -71,11 +72,11 @@ class Grid(geometry.Geometry): def __init__( self, - x: Optional[Sequence[jnp.ndarray]] = None, - grid_size: Optional[Sequence[int]] = None, - cost_fns: Optional[Sequence[costs.CostFn]] = None, - num_a: Optional[int] = None, - grid_dimension: Optional[int] = None, + x: Sequence[jax.Array] | None = None, + grid_size: Sequence[int] | None = None, + cost_fns: Sequence[costs.CostFn] | None = None, + num_a: int | None = None, + grid_dimension: int | None = None, **kwargs: Any, ): super().__init__(**kwargs) @@ -111,7 +112,7 @@ def __init__( } @property - def geometries(self) -> List[geometry.Geometry]: + def geometries(self) -> list[geometry.Geometry]: """Cost matrices along each dimension of the grid.""" geometries = [] for dimension, cost_fn in itertools.zip_longest( @@ -136,7 +137,7 @@ def can_LRC(self) -> bool: # noqa: D102 return True @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self.num_a, self.num_a @property @@ -146,12 +147,12 @@ def is_symmetric(self) -> bool: # noqa: D102 # Reimplemented functions to be used in regularized OT def apply_lse_kernel( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, eps: float, - vec: Optional[jnp.ndarray] = None, + vec: jax.Array | None = None, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: """Apply grid kernel in log space. See notes in parent class for use case. Reshapes vector inputs below as grids, applies kernels onto each slice, and @@ -160,10 +161,10 @@ def apply_lse_kernel( More implementation details in :cite:`schmitz:18`. Args: - f: jnp.ndarray, a vector of potentials - g: jnp.ndarray, a vector of potentials + f: jax.Array, a vector of potentials + g: jax.Array, a vector of potentials eps: float, regularization strength - vec: jnp.ndarray, if needed, a vector onto which apply the kernel weighted + vec: jax.Array, if needed, a vector onto which apply the kernel weighted by f and g. axis: axis (0 or 1) along which summation should be carried out. @@ -210,11 +211,11 @@ def _apply_lse_kernel_one_dimension(self, dimension, f, g, eps, vec=None): def _apply_cost_to_vec( self, - vec: jnp.ndarray, + vec: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + fn: Callable[[jax.Array], jax.Array] | None = None, is_linear: bool = False, - ) -> jnp.ndarray: + ) -> jax.Array: r"""Apply grid's cost matrix (without instantiating it) to a vector. The `apply_cost` operation on grids rests on the following identity. @@ -233,14 +234,14 @@ def _apply_cost_to_vec( summation while keeping dimensions. Args: - vec: jnp.ndarray, flat vector of total size prod(grid_size). + vec: jax.Array, flat vector of total size prod(grid_size). axis: axis 0 if applying transpose costs, 1 if using the original cost. fn: function optionally applied to cost matrix element-wise, before the dot product. is_linear: TODO. Returns: - A jnp.ndarray corresponding to cost x matrix + A jax.Array corresponding to cost x matrix """ # TODO(michalk8): del fn, is_linear @@ -262,10 +263,10 @@ def _apply_cost_to_vec( def apply_kernel( self, - vec: jnp.ndarray, - eps: Optional[float] = None, - axis: Optional[int] = None - ) -> jnp.ndarray: + vec: jax.Array, + eps: float | None = None, + axis: int | None = None + ) -> jax.Array: """Apply grid kernel on scaling vector. See notes in parent class for use. @@ -276,7 +277,7 @@ def apply_kernel( More implementation details in :cite:`schmitz:18`, Args: - vec: jnp.ndarray, a vector of scaling (>0) values. + vec: jax.Array, a vector of scaling (>0) values. eps: float, regularization strength axis: axis (0 or 1) along which summation should be carried out. @@ -294,7 +295,7 @@ def apply_kernel( return vec.ravel() def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 + self, f: jax.Array, g: jax.Array, axis: int = 0 ) -> NoReturn: """Not implemented, use :meth:`apply_transport_from_potentials` instead.""" raise ValueError( @@ -305,7 +306,7 @@ def transport_from_potentials( ) def transport_from_scalings( - self, f: jnp.ndarray, g: jnp.ndarray, axis: int = 0 + self, f: jax.Array, g: jax.Array, axis: int = 0 ) -> NoReturn: """Not implemented, use :meth:`apply_transport_from_scalings` instead.""" raise ValueError( @@ -316,14 +317,14 @@ def transport_from_scalings( ) @property - def cost_matrix(self) -> jnp.ndarray: + def cost_matrix(self) -> jax.Array: """Not implemented.""" raise NotImplementedError( "Instantiating cost matrix is not implemented for grids." ) @property - def kernel_matrix(self) -> jnp.ndarray: + def kernel_matrix(self) -> jax.Array: """Not implemented.""" raise NotImplementedError( "Instantiating kernel matrix is not implemented for grids." @@ -335,7 +336,7 @@ def prepare_divergences( *args: Any, static_b: bool = False, **kwargs: Any - ) -> Tuple["Grid", ...]: + ) -> tuple["Grid", ...]: """Instantiate the geometries used for a divergence computation.""" sep_grid = cls(*args, **kwargs) size = 2 if static_b else 3 diff --git a/src/ott/geometry/low_rank.py b/src/ott/geometry/low_rank.py index 902128b6a..912585242 100644 --- a/src/ott/geometry/low_rank.py +++ b/src/ott/geometry/low_rank.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Literal, Optional, Tuple, Union +from collections.abc import Callable +from typing import Any, Literal import jax import jax.numpy as jnp @@ -48,11 +49,11 @@ class LRCGeometry(geometry.Geometry): def __init__( self, - cost_1: jnp.ndarray, - cost_2: jnp.ndarray, + cost_1: jax.Array, + cost_2: jax.Array, bias: float = 0.0, scale_factor: float = 1.0, - scale_cost: Union[float, Literal["mean", "max_bound", "max_cost"]] = 1.0, + scale_cost: float | Literal["mean", "max_bound", "max_cost"] = 1.0, **kwargs: Any, ): super().__init__(**kwargs) @@ -63,13 +64,13 @@ def __init__( self._scale_cost = scale_cost @property - def cost_1(self) -> jnp.ndarray: + def cost_1(self) -> jax.Array: """First factor of the :attr:`cost_matrix`.""" scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) return scale_factor * self._cost_1 @property - def cost_2(self) -> jnp.ndarray: + def cost_2(self) -> jax.Array: """Second factor of the :attr:`cost_matrix`.""" scale_factor = jnp.sqrt(self._scale_factor * self.inv_scale_cost) return scale_factor * self._cost_2 @@ -84,12 +85,12 @@ def cost_rank(self) -> int: # noqa: D102 return self._cost_1.shape[1] @property - def cost_matrix(self) -> jnp.ndarray: + def cost_matrix(self) -> jax.Array: """Materialize the cost matrix.""" return jnp.matmul(self.cost_1, self.cost_2.T) + self.bias @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self._cost_1.shape[0], self._cost_2.shape[0] @property @@ -98,7 +99,7 @@ def is_symmetric(self) -> bool: # noqa: D102 return (n == m) and jnp.all(self._cost_1 == self._cost_2) @property - def inv_scale_cost(self) -> jnp.ndarray: # noqa: D102 + def inv_scale_cost(self) -> jax.Array: # noqa: D102 if self._scale_cost == "max_bound": x_norm = self._cost_1[:, 0].max() y_norm = self._cost_2[:, 1].max() @@ -117,12 +118,12 @@ def inv_scale_cost(self) -> jnp.ndarray: # noqa: D102 raise ValueError(f"Scaling {self._scale_cost} not implemented.") @property - def diag_cost(self) -> jnp.ndarray: + def diag_cost(self) -> jax.Array: """Diagonal of the cost matrix.""" assert self.is_square, "Diagonal cost only available for square geometries." return jnp.sum(self._cost_1 * self._cost_2, axis=-1) - def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + def apply_square_cost(self, arr: jax.Array, axis: int = 0) -> jax.Array: """Apply elementwise-square of cost matrix to array (vector or matrix).""" (n, m), r = self.shape, self.cost_rank # When applying square of a LRCGeometry, one can either elementwise square @@ -140,15 +141,15 @@ def apply_square_cost(self, arr: jnp.ndarray, axis: int = 0) -> jnp.ndarray: def _apply_cost_to_vec( self, - vec: jnp.ndarray, + vec: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + fn: Callable[[jax.Array], jax.Array] | None = None, is_linear: bool = False, - ) -> jnp.ndarray: + ) -> jax.Array: """Apply [num_a, num_b] fn(cost) (or transpose) to vector. Args: - vec: jnp.ndarray [num_a,] ([num_b,] if axis=1) vector + vec: jax.Array [num_a,] ([num_b,] if axis=1) vector axis: axis on which the reduction is done. fn: function optionally applied to cost matrix element-wise, before the doc product @@ -156,7 +157,7 @@ def _apply_cost_to_vec( implementation. Returns: - A jnp.ndarray corresponding to cost x vector + A jax.Array corresponding to cost x vector """ if fn is None or is_linear: return self._apply_cost_to_vec_fast(vec, axis, fn=fn) @@ -164,10 +165,10 @@ def _apply_cost_to_vec( def _apply_cost_to_vec_fast( self, - vec: jnp.ndarray, + vec: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - ) -> jnp.ndarray: + fn: Callable[[jax.Array], jax.Array] | None = None, + ) -> jax.Array: c1, c2 = (self.cost_1, self.cost_2) if axis == 1 else (self.cost_2, self.cost_1) bias = self.bias @@ -177,7 +178,7 @@ def _apply_cost_to_vec_fast( return out + bias * jnp.sum(vec) * jnp.ones_like(out) @property - def _max_cost_matrix(self) -> jnp.ndarray: + def _max_cost_matrix(self) -> jax.Array: fn = utils.batched_vmap( lambda c1, c2: jnp.max(c1 @ c2.T), batch_size=1024, in_axes=(0, None) ) @@ -187,7 +188,7 @@ def to_LRCGeometry( self, rank: int = 0, tol: float = 1e-2, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, scale: float = 1.0, ) -> "LRCGeometry": """Return self.""" @@ -256,9 +257,9 @@ class LRKGeometry(geometry.Geometry): def __init__( self, - k1: jnp.ndarray, - k2: jnp.ndarray, - epsilon: Optional[float] = None, + k1: jax.Array, + k2: jax.Array, + epsilon: float | None = None, **kwargs: Any ): super().__init__(epsilon=epsilon, relative_epsilon=None, **kwargs) @@ -268,14 +269,14 @@ def __init__( @classmethod def from_pointcloud( cls, - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, *, kernel: Literal["gaussian", "arccos"], rank: int = 100, std: float = 1.0, n: int = 1, - rng: Optional[jax.Array] = None + rng: jax.Array | None = None ) -> "LRKGeometry": r"""Low-rank kernel approximation :cite:`scetbon:20`. @@ -316,20 +317,20 @@ def from_pointcloud( def apply_kernel( # noqa: D102 self, - vec: jnp.ndarray, - eps: Optional[float] = None, + vec: jax.Array, + eps: float | None = None, axis: int = 0, - ) -> jnp.ndarray: + ) -> jax.Array: if axis == 0: return self.k2 @ (self.k1.T @ vec) return self.k1 @ (self.k2.T @ vec) @property - def kernel_matrix(self) -> jnp.ndarray: # noqa: D102 + def kernel_matrix(self) -> jax.Array: # noqa: D102 return self.k1 @ self.k2.T @property - def cost_matrix(self) -> jnp.ndarray: # noqa: D102 + def cost_matrix(self) -> jax.Array: # noqa: D102 eps = jnp.finfo(self.dtype).tiny return -self.epsilon * jnp.log(self.kernel_matrix + eps) @@ -338,16 +339,14 @@ def rank(self) -> int: # noqa: D102 return self.k1.shape[1] @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self.k1.shape[0], self.k2.shape[0] @property def dtype(self) -> jnp.dtype: # noqa: D102 return self.k1.dtype - def transport_from_potentials( - self, f: jnp.ndarray, g: jnp.ndarray - ) -> jnp.ndarray: + def transport_from_potentials(self, f: jax.Array, g: jax.Array) -> jax.Array: """Not implemented.""" raise ValueError("Not implemented.") @@ -361,11 +360,11 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 def _gaussian_kernel( rng: jax.Array, - x: jnp.ndarray, + x: jax.Array, n_features: int, eps: float, - R: jnp.ndarray, -) -> jnp.ndarray: + R: jax.Array, +) -> jax.Array: _, d = x.shape cost_fn = costs.SqEuclidean() @@ -385,12 +384,12 @@ def _gaussian_kernel( def _arccos_kernel( rng: jax.Array, - x: jnp.ndarray, + x: jax.Array, n_features: int, n: int, std: float = 1.0, kappa: float = 1e-6, -) -> jnp.ndarray: +) -> jax.Array: n_points, d = x.shape c = jnp.sqrt(2) * (std ** (d / 2)) diff --git a/src/ott/geometry/pointcloud.py b/src/ott/geometry/pointcloud.py index 37785a231..433af4e75 100644 --- a/src/ott/geometry/pointcloud.py +++ b/src/ott/geometry/pointcloud.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Literal, Optional, Tuple, Union +from collections.abc import Callable +from typing import Any, Literal import jax import jax.numpy as jnp @@ -53,12 +54,12 @@ class PointCloud(geometry.Geometry): def __init__( self, - x: jnp.ndarray, - y: Optional[jnp.ndarray] = None, - cost_fn: Optional[costs.CostFn] = None, - batch_size: Optional[int] = None, - scale_cost: Union[float, Literal["mean", "max_norm", "max_bound", - "max_cost", "median"]] = 1.0, + x: jax.Array, + y: jax.Array | None = None, + cost_fn: costs.CostFn | None = None, + batch_size: int | None = None, + scale_cost: float + | Literal["mean", "max_norm", "max_bound", "max_cost", "median"] = 1.0, **kwargs: Any, ): super().__init__(**kwargs) @@ -73,17 +74,17 @@ def __init__( def apply_lse_kernel( # noqa: D102 self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, eps: float, - vec: Optional[jnp.ndarray] = None, + vec: jax.Array | None = None, axis: int = 0 - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array]: if not self.is_online: return super().apply_lse_kernel(f, g, eps, vec, axis) - def apply(x: jnp.ndarray, y: jnp.ndarray, f: jnp.ndarray, - g: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + def apply(x: jax.Array, y: jax.Array, f: jax.Array, + g: jax.Array) -> tuple[jax.Array, jax.Array]: x, y = jnp.atleast_2d(x), jnp.atleast_2d(y) cost = self.cost_fn.all_pairs(x, y) * inv_scale_cost cost = cost.squeeze(1 - axis) @@ -104,16 +105,16 @@ def apply(x: jnp.ndarray, y: jnp.ndarray, f: jnp.ndarray, def apply_kernel( # noqa: D102 self, - vec: jnp.ndarray, - eps: Optional[float] = None, + vec: jax.Array, + eps: float | None = None, axis: int = 0 - ) -> jnp.ndarray: + ) -> jax.Array: if eps is None: eps = self.epsilon if not self.is_online: return super().apply_kernel(vec, eps, axis) - def apply(x: jnp.ndarray, y: jnp.ndarray, vec: jnp.ndarray) -> jnp.ndarray: + def apply(x: jax.Array, y: jax.Array, vec: jax.Array) -> jax.Array: x, y = jnp.atleast_2d(x), jnp.atleast_2d(y) cost = self.cost_fn.all_pairs(x, y) * inv_scale_cost cost = cost.squeeze(1 - axis) @@ -128,14 +129,14 @@ def apply(x: jnp.ndarray, y: jnp.ndarray, vec: jnp.ndarray) -> jnp.ndarray: def _apply_cost_to_vec( self, - vec: jnp.ndarray, + vec: jax.Array, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + fn: Callable[[jax.Array], jax.Array] | None = None, is_linear: bool = False, - scale_cost: Optional[float] = None, - ) -> jnp.ndarray: + scale_cost: float | None = None, + ) -> jax.Array: - def apply(x: jnp.ndarray, y: jnp.ndarray, arr: jnp.ndarray) -> jnp.ndarray: + def apply(x: jax.Array, y: jax.Array, arr: jax.Array) -> jax.Array: x, y = jnp.atleast_2d(x), jnp.atleast_2d(y) cost = self.cost_fn.all_pairs(x, y) * scale_cost cost = cost.squeeze(1 - axis) @@ -169,11 +170,11 @@ def apply(x: jnp.ndarray, y: jnp.ndarray, arr: jnp.ndarray) -> jnp.ndarray: def _apply_sqeucl_cost( self, - vec: jnp.ndarray, + vec: jax.Array, scale_cost: float, axis: int = 0, - fn: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, - ) -> jnp.ndarray: + fn: Callable[[jax.Array], jax.Array] | None = None, + ) -> jax.Array: assert vec.ndim == 1, vec.shape assert self.is_squared_euclidean, "Cost matrix is not a squared Euclidean." x, y = (self.x, self.y) if axis == 0 else (self.y, self.x) @@ -187,7 +188,7 @@ def _apply_sqeucl_cost( def _compute_summary_online( self, summary: Literal["mean", "max_cost"] - ) -> jnp.ndarray: + ) -> jax.Array: """Compute mean or max of cost matrix online, i.e. without instantiating it. Args: @@ -197,7 +198,7 @@ def _compute_summary_online( summary statistics """ - def compute_max(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + def compute_max(x: jax.Array, y: jax.Array) -> jax.Array: x, y = jnp.atleast_2d(x), jnp.atleast_2d(y) cost = self.cost_fn.all_pairs(x, y) return jnp.max(jnp.abs(cost)) @@ -218,18 +219,18 @@ def compute_max(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: f"Scaling method {summary} does not exist for online mode." ) - def barycenter(self, weights: jnp.ndarray) -> jnp.ndarray: + def barycenter(self, weights: jax.Array) -> jax.Array: """Compute barycenter of points in self.x using weights.""" return self.cost_fn.barycenter(self.x, weights)[0] @classmethod def prepare_divergences( cls, - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, static_b: bool = False, **kwargs: Any - ) -> Tuple["PointCloud", ...]: + ) -> tuple["PointCloud", ...]: """Instantiate the geometries used for a divergence computation.""" couples = [(x, y), (x, x)] if not static_b: @@ -267,7 +268,7 @@ def to_LRCGeometry( self, scale: float = 1.0, **kwargs: Any, - ) -> Union[low_rank.LRCGeometry, "PointCloud"]: + ) -> "low_rank.LRCGeometry | PointCloud": r"""Convert point cloud to low-rank geometry. Args: @@ -332,15 +333,15 @@ def _dotp_to_lr(self, scale: float = 1.0) -> low_rank.LRCGeometry: ) @property - def cost_matrix(self) -> Optional[jnp.ndarray]: # noqa: D102 + def cost_matrix(self) -> jax.Array | None: # noqa: D102 return self.inv_scale_cost * self._unscaled_cost_matrix @property - def _unscaled_cost_matrix(self) -> jnp.ndarray: + def _unscaled_cost_matrix(self) -> jax.Array: return self.cost_fn.all_pairs(self.x, self.y) @property - def inv_scale_cost(self) -> jnp.ndarray: # noqa: D102 + def inv_scale_cost(self) -> jax.Array: # noqa: D102 if self._scale_cost == "max_cost": if self.is_online: return 1.0 / self._compute_summary_online(self._scale_cost) @@ -384,8 +385,8 @@ def inv_scale_cost(self) -> jnp.ndarray: # noqa: D102 def subset( # noqa: D102 self, - row_ixs: Optional[jnp.ndarray] = None, - col_ixs: Optional[jnp.ndarray] = None, + row_ixs: jax.Array | None = None, + col_ixs: jax.Array | None = None, ) -> "PointCloud": (x, y, *rest), aux_data = self.tree_flatten() if row_ixs is not None: @@ -395,11 +396,11 @@ def subset( # noqa: D102 return type(self).tree_unflatten(aux_data, (x, y, *rest)) @property - def kernel_matrix(self) -> Optional[jnp.ndarray]: # noqa: D102 + def kernel_matrix(self) -> jax.Array | None: # noqa: D102 return jnp.exp(-self.cost_matrix / self.epsilon) @property - def shape(self) -> Tuple[int, int]: # noqa: D102 + def shape(self) -> tuple[int, int]: # noqa: D102 return self.x.shape[0], self.y.shape[0] @property @@ -435,7 +436,7 @@ def cost_rank(self) -> int: # noqa: D102 return self.x.shape[1] @property - def batch_size(self) -> Optional[int]: + def batch_size(self) -> int | None: """Batch size for online mode.""" if self._batch_size is None: return None @@ -448,7 +449,7 @@ def is_online(self) -> bool: return self.batch_size is not None @property - def diag_cost(self) -> jnp.ndarray: + def diag_cost(self) -> jax.Array: """Diagonal of the cost matrix.""" assert self.is_square, "Cost matrix must be square to compute diagonal." return jax.vmap(self.cost_fn, in_axes=(0, 0))(self.x, self.y) diff --git a/src/ott/geometry/regularizers.py b/src/ott/geometry/regularizers.py index 9d6cb56b3..80b21a12d 100644 --- a/src/ott/geometry/regularizers.py +++ b/src/ott/geometry/regularizers.py @@ -13,7 +13,8 @@ # limitations under the License. import abc import functools -from typing import Any, Callable, Optional, Tuple, Union +from collections.abc import Callable +from typing import Any import lineax as lx @@ -37,7 +38,7 @@ class ProximalOperator(abc.ABC): """Proximal operator base class.""" @abc.abstractmethod - def __call__(self, x: jnp.ndarray) -> float: + def __call__(self, x: jax.Array) -> float: """Function. Args: @@ -48,7 +49,7 @@ def __call__(self, x: jnp.ndarray) -> float: """ @abc.abstractmethod - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: """Proximal operator. Args: @@ -59,7 +60,7 @@ def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: The prox of ``v``. """ - def prox_dual(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + def prox_dual(self, v: jax.Array, tau: float = 1.0) -> jax.Array: r"""Proximal operator of the convex conjugate. Uses Moreau's decomposition: @@ -77,7 +78,7 @@ def prox_dual(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: """ return v - tau * self.prox(v / tau, 1.0 / tau) - def moreau_envelope(self, x: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: + def moreau_envelope(self, x: jax.Array, tau: float = 1.0) -> jax.Array: r"""Moreau Envelope. Uses Remark 12.24 from :cite:`bauschke:17`: @@ -120,10 +121,10 @@ def __init__(self, f: ProximalOperator, alpha: float = 1.0, b: float = 0.0): self.alpha = alpha self.b = b - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 return self.alpha * self.f(x) + self.b - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 return self.f.prox(v, tau * self.alpha) def tree_flatten(self): # noqa: D102 @@ -143,7 +144,7 @@ class Regularization(ProximalOperator): def __init__( self, f: ProximalOperator, - a: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, rho: float = 1.0, ): super().__init__() @@ -151,11 +152,11 @@ def __init__( self.a = a self.rho = rho - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 norm = jnp.sum(x ** 2) if self.a is None else jnp.sum((x - self.a) ** 2) return self.f(x) + (0.5 * self.rho) * norm - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 tau_tilde = tau / (1.0 + tau * self.rho) # (tau_tilde / tau) * v vv = 1.0 / (1 + tau * self.rho) * v @@ -185,25 +186,25 @@ class Orthogonal(ProximalOperator): def __init__( self, f: ProximalOperator, - A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]], - b: Optional[jnp.ndarray] = None, + A: jax.Array | lx.AbstractLinearOperator | None, + b: jax.Array | None = None, nu: float = 1.0, ): assert nu > 0.0, nu super().__init__() self.f = f # AA^T = alpha I - self.A = lx.MatrixLinearOperator(A) if isinstance(A, jnp.ndarray) else A + self.A = lx.MatrixLinearOperator(A) if isinstance(A, jax.Array) else A self.b = b self.nu = nu - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 z = self.A.mv(x) if self.b is not None: z = z + self.b return self.f(z) - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 w = self.A.mv(v) if self.b is None: tmp = self.f.prox(w, tau * self.nu) @@ -244,29 +245,29 @@ class Quadratic(ProximalOperator): def __init__( self, - A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]] = None, - b: Optional[jnp.ndarray] = None, + A: jax.Array | lx.AbstractLinearOperator | None = None, + b: jax.Array | None = None, *, is_complement: bool = False, is_orthogonal: bool = False, is_factor: bool = False, - solver: Optional[Callable[[lx.AbstractLinearOperator, jnp.ndarray], - jnp.ndarray]] = None, + solver: Callable[[lx.AbstractLinearOperator, jax.Array], jax.Array] + | None = None, ): super().__init__() - self.A = lx.MatrixLinearOperator(A) if isinstance(A, jnp.ndarray) else A + self.A = lx.MatrixLinearOperator(A) if isinstance(A, jax.Array) else A self.b = b self._is_complement = is_complement self._is_orthogonal = is_orthogonal self._is_factor = is_factor self._solver = solver - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 Q = self.Q y = 0.5 * (jnp.dot(x, x) if Q is None else jnp.dot(x, Q.mv(x))) return y if self.b is None else (y + jnp.dot(x, self.b)) - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 # section 6.1.1 in :cite:`parikh:14` Q = self.Q b = v if self.b is None else (v - tau * self.b) @@ -295,7 +296,7 @@ def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 return self._solver(A, b) @property - def A_comp(self) -> Optional[lx.AbstractLinearOperator]: + def A_comp(self) -> lx.AbstractLinearOperator | None: r"""Orthogonal complement :math:`A^{\perp}` of :math:`A`.""" return _complement( self.A, self.is_orthogonal @@ -317,7 +318,7 @@ def is_orthogonal(self) -> bool: return self.A is not None and self._is_orthogonal @property - def Q(self) -> Optional[lx.AbstractLinearOperator]: + def Q(self) -> lx.AbstractLinearOperator | None: r"""Linear operator :math:`Q`.""" Q = self.A_comp if self.is_complement else self.A if Q is None: @@ -337,10 +338,10 @@ def tree_flatten(self): # noqa: D102 class L1(ProximalOperator): r"""L1-norm regularizer :math:`\ell_1`.""" - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 return jnp.linalg.norm(x, ord=1) - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 return jnp.sign(v) * jax.nn.relu(jnp.abs(v) - tau) @@ -356,17 +357,17 @@ class SqL2(ProximalOperator): def __init__( self, - A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]] = None, + A: jax.Array | lx.AbstractLinearOperator | None = None, **kwargs: Any, ): super().__init__() self.f = Quadratic(A, is_factor=True, **kwargs) self._init_kwargs = kwargs - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 return self.f(x) - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 return self.f.prox(v, tau) def tree_flatten(self): # noqa: D102 @@ -393,13 +394,13 @@ def __init__(self, gamma: float = 1.0): super().__init__() self.gamma = gamma - def __call__(self, x: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, x: jax.Array) -> float: # noqa: D102 # Lemma 2.1 of `schreck:15` u = jnp.arcsinh(jnp.abs(x) / (2.0 * self.gamma)) y = u - 0.5 * jnp.exp(-2.0 * u) return self.gamma ** 2 * jnp.sum(y + 0.5) # make positive - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> jax.Array: # noqa: D102 s = (tau * self.gamma) ** 2 return jnp.where(v ** 2 <= s, 0.0, v - s / jnp.where(v == 0.0, 1.0, v)) @@ -428,7 +429,7 @@ def __init__(self, k: int): super().__init__() self.k = k - def __call__(self, z: jnp.ndarray) -> float: # noqa: D102 + def __call__(self, z: jax.Array) -> float: # noqa: D102 # Prop 2.1 in :cite:`argyriou:12` k = self.k top_w = jax.lax.top_k(jnp.abs(z), k)[0] # Fetch largest k values @@ -449,15 +450,14 @@ def __call__(self, z: jnp.ndarray) -> float: # noqa: D102 return 0.5 * (s + (r + 1) * cesaro[r] ** 2) - def prox(self, v: jnp.ndarray, tau: float = 1.0) -> float: # noqa: D102 + def prox(self, v: jax.Array, tau: float = 1.0) -> float: # noqa: D102 @functools.partial(jax.vmap, in_axes=[0, None, None]) - def find_indices(r: int, l: jnp.ndarray, - z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + def find_indices(r: int, l: jax.Array, + z: jax.Array) -> tuple[jax.Array, jax.Array]: @functools.partial(jax.vmap, in_axes=[None, 0, None]) - def inner(r: int, l: int, - z: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + def inner(r: int, l: int, z: jax.Array) -> tuple[jax.Array, jax.Array]: i = k - r - 1 res = jnp.sum(z * ((i <= ixs) & (ixs < l))) res /= l - k + (beta + 1) * r + beta + 1 diff --git a/src/ott/geometry/segment.py b/src/ott/geometry/segment.py index 34b7efcce..bfd916b1e 100644 --- a/src/ott/geometry/segment.py +++ b/src/ott/geometry/segment.py @@ -12,7 +12,7 @@ # 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 typing import Callable, Optional, Tuple +from collections.abc import Callable import jax import jax.numpy as jnp @@ -21,15 +21,15 @@ def segment_point_cloud( - x: jnp.ndarray, - a: Optional[jnp.ndarray] = None, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - segment_ids: Optional[jnp.ndarray] = None, + x: jax.Array, + a: jax.Array | None = None, + num_segments: int | None = None, + max_measure_size: int | None = None, + segment_ids: jax.Array | None = None, indices_are_sorted: bool = False, - num_per_segment: Optional[Tuple[int, ...]] = None, - padding_vector: Optional[jnp.ndarray] = None -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + num_per_segment: tuple[int, ...] | None = None, + padding_vector: jax.Array | None = None +) -> tuple[jax.Array, jax.Array, jax.Array]: """Segment and pad as needed the entries of a point cloud. There are two interfaces: @@ -131,21 +131,20 @@ def segment_point_cloud( def _segment_interface( - x: jnp.ndarray, - y: jnp.ndarray, - eval_fn: Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray], - jnp.ndarray], - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, + x: jax.Array, + y: jax.Array, + eval_fn: Callable[[jax.Array, jax.Array, jax.Array, jax.Array], jax.Array], + num_segments: int | None = None, + max_measure_size: int | None = None, + segment_ids_x: jax.Array | None = None, + segment_ids_y: jax.Array | None = None, indices_are_sorted: bool = False, - num_per_segment_x: Optional[jnp.ndarray] = None, - num_per_segment_y: Optional[jnp.ndarray] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, - padding_vector: Optional[jnp.ndarray] = None, -) -> jnp.ndarray: + num_per_segment_x: jax.Array | None = None, + num_per_segment_y: jax.Array | None = None, + weights_x: jax.Array | None = None, + weights_y: jax.Array | None = None, + padding_vector: jax.Array | None = None, +) -> jax.Array: """Wrapper to segment two point clouds and return parallel evaluations. Utility function that segments two point clouds using the approach outlined diff --git a/src/ott/geometry/semidiscrete_pointcloud.py b/src/ott/geometry/semidiscrete_pointcloud.py index 421122bee..21fa786e4 100644 --- a/src/ott/geometry/semidiscrete_pointcloud.py +++ b/src/ott/geometry/semidiscrete_pointcloud.py @@ -11,12 +11,14 @@ # 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 typing import Callable, Literal, Optional, Tuple, Union +from collections.abc import Callable +from typing import Literal import jax import jax.numpy as jnp import jax.random as jr import jax.tree_util as jtu +from jax.typing import ArrayLike, DTypeLike from ott.geometry import costs, pointcloud @@ -48,14 +50,14 @@ class SemidiscretePointCloud: def __init__( self, - sampler: Callable[[jax.Array, Tuple[int, ...], Optional[jnp.dtype]], + sampler: Callable[[jax.Array, tuple[int, ...], DTypeLike | None], jax.Array], y: jax.Array, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[Union[float, jax.Array]] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_norm", "max_bound", - "max_cost", "median"]] = 1.0, + cost_fn: costs.CostFn | None = None, + epsilon: ArrayLike | None = None, + relative_epsilon: Literal["mean", "std"] | None = None, + scale_cost: float + | Literal["mean", "max_norm", "max_bound", "max_cost", "median"] = 1.0, relative_epsilon_seed: int = 0, relative_epsilon_num_samples: int = 1024, ): @@ -75,7 +77,7 @@ def sample( rng: jax.Array, num_samples: int, *, - epsilon: Optional[float] = None + epsilon: float | None = None ) -> pointcloud.PointCloud: """Sample a point cloud. diff --git a/src/ott/initializers/linear/initializers.py b/src/ott/initializers/linear/initializers.py index daa70ae14..a03ae1f01 100644 --- a/src/ott/initializers/linear/initializers.py +++ b/src/ott/initializers/linear/initializers.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc -from typing import Any, Dict, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import jax import jax.numpy as jnp @@ -36,8 +37,8 @@ def init_fu( self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: """Initialize Sinkhorn potential/scaling f_u. Args: @@ -54,8 +55,8 @@ def init_gv( self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: """Initialize Sinkhorn potential/scaling g_v. Args: @@ -71,8 +72,8 @@ def __call__( self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + rng: jax.Array | None = None, + ) -> tuple[jax.Array, jax.Array]: """Initialize Sinkhorn potentials/scalings f_u and g_v. Args: @@ -98,12 +99,12 @@ def __call__( gv = jnp.where(ot_prob.b > 0.0, gv, mask_value) return fu, gv - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [], {} @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "SinkhornInitializer": return cls(*children, **aux_data) @@ -116,8 +117,8 @@ def init_fu( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: del rng return jnp.zeros_like(ot_prob.a) if lse_mode else jnp.ones_like(ot_prob.a) @@ -125,8 +126,8 @@ def init_gv( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: del rng return jnp.zeros_like(ot_prob.b) if lse_mode else jnp.ones_like(ot_prob.b) @@ -146,8 +147,8 @@ def init_fu( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: # import Gaussian here due to circular imports from ott.tools.gaussian_mixture import gaussian @@ -195,8 +196,8 @@ def __init__( self.vectorized_update = vectorized_update def _init_sorting_dual( - self, modified_cost: jnp.ndarray, init_f: jnp.ndarray - ) -> jnp.ndarray: + self, modified_cost: jax.Array, init_f: jax.Array + ) -> jax.Array: """Run DualSort algorithm. Args: @@ -209,15 +210,15 @@ def _init_sorting_dual( """ def body_fn( - state: Tuple[jnp.ndarray, float, int] - ) -> Tuple[jnp.ndarray, float, int]: + state: tuple[jax.Array, float, int] + ) -> tuple[jax.Array, float, int]: prev_f, _, it = state new_f = fn(prev_f, modified_cost) diff = jnp.sum((new_f - prev_f) ** 2) it += 1 return new_f, diff, it - def cond_fn(state: Tuple[jnp.ndarray, float, int]) -> bool: + def cond_fn(state: tuple[jax.Array, float, int]) -> bool: _, diff, it = state return jnp.logical_and(diff > self.tolerance, it < self.max_iter) @@ -233,9 +234,9 @@ def init_fu( self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - init_f: Optional[jnp.ndarray] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + init_f: jax.Array | None = None, + ) -> jax.Array: """Apply DualSort algorithm. Args: @@ -270,7 +271,7 @@ def init_fu( f_potential ) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([], { "tolerance": self.tolerance, "max_iter": self.max_iter, @@ -299,7 +300,7 @@ class SubsampleInitializer(DefaultInitializer): def __init__( self, subsample_n_x: int, - subsample_n_y: Optional[int] = None, + subsample_n_y: int | None = None, **kwargs: Any, ): super().__init__() @@ -311,8 +312,8 @@ def init_fu( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: from ott.solvers import linear assert isinstance( @@ -352,7 +353,7 @@ def init_fu( # noqa: D102 f_potential ) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([], { "subsample_n_x": self.subsample_n_x, "subsample_n_y": self.subsample_n_y, @@ -360,9 +361,7 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 }) -def _vectorized_update( - f: jnp.ndarray, modified_cost: jnp.ndarray -) -> jnp.ndarray: +def _vectorized_update(f: jax.Array, modified_cost: jax.Array) -> jax.Array: """Inner loop DualSort Update. Args: @@ -375,9 +374,7 @@ def _vectorized_update( return jnp.min(modified_cost + f[None, :], axis=1) -def _coordinate_update( - f: jnp.ndarray, modified_cost: jnp.ndarray -) -> jnp.ndarray: +def _coordinate_update(f: jax.Array, modified_cost: jax.Array) -> jax.Array: """Coordinate-wise updates within inner loop. Args: @@ -388,7 +385,7 @@ def _coordinate_update( updated potential vector, f. """ - def body_fn(i: int, f: jnp.ndarray) -> jnp.ndarray: + def body_fn(i: int, f: jax.Array) -> jax.Array: new_f = jnp.min(modified_cost[i, :] + f) return f.at[i].set(new_f) diff --git a/src/ott/initializers/linear/initializers_lr.py b/src/ott/initializers/linear/initializers_lr.py index d939853db..71a11c987 100644 --- a/src/ott/initializers/linear/initializers_lr.py +++ b/src/ott/initializers/linear/initializers_lr.py @@ -13,18 +13,8 @@ # limitations under the License. import abc import functools -from typing import ( - TYPE_CHECKING, - Any, - Dict, - Literal, - Mapping, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Union import jax import jax.numpy as jnp @@ -68,9 +58,9 @@ def init_q( ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: """Initialize the low-rank factor :math:`Q`. Args: @@ -89,9 +79,9 @@ def init_r( ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: """Initialize the low-rank factor :math:`R`. Args: @@ -110,7 +100,7 @@ def init_g( ot_prob: Problem_t, rng: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: """Initialize the low-rank factor :math:`g`. Args: @@ -125,9 +115,9 @@ def init_g( def __call__( self, ot_prob: Problem_t, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, **kwargs: Any, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Initialize the factors :math:`Q`, :math:`R` and :math:`g`. Args: @@ -157,12 +147,12 @@ def rank(self) -> int: """Rank of the transport matrix factorization.""" return self._rank - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [], {**self._kwargs, "rank": self.rank} @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "LRInitializer": return cls(*children, **aux_data) @@ -181,9 +171,9 @@ def init_q( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del kwargs, init_g a = ot_prob.a init_q = jnp.abs(jax.random.normal(rng, (a.shape[0], self.rank))) @@ -194,9 +184,9 @@ def init_r( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del kwargs, init_g b = ot_prob.b init_r = jnp.abs(jax.random.normal(rng, (b.shape[0], self.rank))) @@ -207,7 +197,7 @@ def init_g( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del kwargs init_g = jnp.abs(jax.random.uniform(rng, (self.rank,))) + 1.0 return init_g / jnp.sum(init_g) @@ -225,10 +215,10 @@ class Rank2Initializer(LRInitializer): def _compute_factor( self, ot_prob: Problem_t, - init_g: jnp.ndarray, + init_g: jax.Array, *, which: Literal["q", "r"], - ) -> jnp.ndarray: + ) -> jax.Array: a, b = ot_prob.a, ot_prob.b marginal = a if which == "q" else b n, r = marginal.shape[0], self.rank @@ -254,9 +244,9 @@ def init_q( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del rng, kwargs return self._compute_factor(ot_prob, init_g, which="q") @@ -265,9 +255,9 @@ def init_r( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del rng, kwargs return self._compute_factor(ot_prob, init_g, which="r") @@ -276,7 +266,7 @@ def init_g( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del rng, kwargs return jnp.ones((self.rank,)) / self.rank @@ -302,7 +292,7 @@ def __init__( rank: int, min_iterations: int = 100, max_iterations: int = 100, - sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, + sinkhorn_kwargs: Mapping[str, Any] | None = None, **kwargs: Any ): super().__init__(rank, **kwargs) @@ -311,7 +301,7 @@ def __init__( self._sinkhorn_kwargs = {} if sinkhorn_kwargs is None else sinkhorn_kwargs @staticmethod - def _extract_array(geom: geometry.Geometry, *, first: bool) -> jnp.ndarray: + def _extract_array(geom: geometry.Geometry, *, first: bool) -> jax.Array: if isinstance(geom, pointcloud.PointCloud): return geom.x if first else geom.y if isinstance(geom, low_rank.LRCGeometry): @@ -325,10 +315,10 @@ def _compute_factor( ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, which: Literal["q", "r"], **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: from ott.problems.linear import linear_problem from ott.problems.quadratic import quadratic_problem from ott.solvers.linear import sinkhorn @@ -367,9 +357,9 @@ def init_q( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: return self._compute_factor( ot_prob, rng, init_g=init_g, which="q", **kwargs ) @@ -379,9 +369,9 @@ def init_r( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: return self._compute_factor( ot_prob, rng, init_g=init_g, which="r", **kwargs ) @@ -391,11 +381,11 @@ def init_g( # noqa: D102 ot_prob: Problem_t, rng: jax.Array, **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: del rng, kwargs return jnp.ones((self.rank,)) / self.rank - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 children, aux_data = super().tree_flatten() aux_data["sinkhorn_kwargs"] = self._sinkhorn_kwargs aux_data["min_iterations"] = self._min_iter @@ -429,7 +419,7 @@ def __init__( max_iterations: int = 100, inner_iterations: int = 10, threshold: float = 1e-6, - sinkhorn_kwargs: Optional[Mapping[str, Any]] = None, + sinkhorn_kwargs: Mapping[str, Any] | None = None, ): super().__init__( rank, @@ -445,14 +435,14 @@ def __init__( class Constants(NamedTuple): # noqa: D106 solver: "sinkhorn.Sinkhorn" geom: geometry.Geometry # (n, n) - marginal: jnp.ndarray # (n,) - g: jnp.ndarray # (r,) + marginal: jax.Array # (n,) + g: jax.Array # (r,) gamma: float threshold: float class State(NamedTuple): # noqa: D106 - factor: jnp.ndarray - criterions: jnp.ndarray + factor: jax.Array + criterions: jax.Array crossed_threshold: bool def _compute_factor( @@ -460,10 +450,10 @@ def _compute_factor( ot_prob: Problem_t, rng: jax.Array, *, - init_g: jnp.ndarray, + init_g: jax.Array, which: Literal["q", "r"], **kwargs: Any, - ) -> jnp.ndarray: + ) -> jax.Array: from ott.problems.linear import linear_problem from ott.problems.quadratic import quadratic_problem from ott.solvers.linear import sinkhorn diff --git a/src/ott/initializers/neural/meta_initializer.py b/src/ott/initializers/neural/meta_initializer.py index f13539e03..a8d84396a 100644 --- a/src/ott/initializers/neural/meta_initializer.py +++ b/src/ott/initializers/neural/meta_initializer.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Dict, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import jax import jax.numpy as jnp @@ -74,10 +75,9 @@ def __init__( self, geom: geometry.Geometry, meta_model: nn.Module, - opt: Optional[optax.GradientTransformation - ] = optax.adam(learning_rate=1e-3), # noqa: B008 - rng: Optional[jax.Array] = None, - state: Optional[train_state.TrainState] = None, + opt: optax.GradientTransformation | None = optax.adam(1e-3), # noqa: B008 + rng: jax.Array | None = None, + state: train_state.TrainState | None = None, ): self.geom = geom self.opt = opt @@ -99,8 +99,8 @@ def __init__( self.update_impl = self._get_update_fn() def update( - self, state: train_state.TrainState, a: jnp.ndarray, b: jnp.ndarray - ) -> Tuple[jnp.ndarray, jnp.ndarray, train_state.TrainState]: + self, state: train_state.TrainState, a: jax.Array, b: jax.Array + ) -> tuple[jax.Array, jax.Array, train_state.TrainState]: r"""Update the meta model with the dual objective. The goal is for the model to match the optimal duals, i.e., @@ -138,8 +138,8 @@ def init_fu( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool, - rng: Optional[jax.Array] = None, - ) -> jnp.ndarray: + rng: jax.Array | None = None, + ) -> jax.Array: del rng # Detect if the problem is batched. assert ot_prob.a.ndim in (1, 2) @@ -190,9 +190,9 @@ def update(state, a, b): return update def _compute_f( - self, a: jnp.ndarray, b: jnp.ndarray, - params: frozen_dict.FrozenDict[str, jnp.ndarray] - ) -> jnp.ndarray: + self, a: jax.Array, b: jax.Array, + params: frozen_dict.FrozenDict[str, jax.Array] + ) -> jax.Array: r"""Predict the optimal :math:`f` potential. Args: @@ -206,7 +206,7 @@ def _compute_f( return self.meta_model.apply({"params": params}, jnp.concatenate([a, b], axis=-1)) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [self.geom, self.meta_model, self.opt], { "rng": self.rng, "state": self.state diff --git a/src/ott/initializers/quadratic/initializers.py b/src/ott/initializers/quadratic/initializers.py index 373f4d56c..d8d36173d 100644 --- a/src/ott/initializers/quadratic/initializers.py +++ b/src/ott/initializers/quadratic/initializers.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc -from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Literal +import jax import jax.numpy as jnp import jax.tree_util as jtu @@ -72,12 +74,12 @@ def _create_geometry( Geometry used to initialize the linearized problem. """ - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [], {} @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "BaseQuadraticInitializer": return cls(*children, **aux_data) @@ -119,7 +121,7 @@ class QuadraticInitializer(BaseQuadraticInitializer): defaults to the product coupling :math:`ab^T`. """ - def __init__(self, init_coupling: Optional[jnp.ndarray] = None): + def __init__(self, init_coupling: jax.Array | None = None): super().__init__() self.init_coupling = init_coupling @@ -128,7 +130,7 @@ def _create_geometry( quad_prob: "quadratic_problem.QuadraticProblem", *, epsilon: float, - relative_epsilon: Optional[Literal["mean", "std"]] = None, + relative_epsilon: Literal["mean", "std"] | None = None, **kwargs: Any, ) -> geometry.Geometry: """Compute initial geometry for linearization. @@ -179,5 +181,5 @@ def _create_geometry( relative_epsilon=relative_epsilon ) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return [self.init_coupling], {} diff --git a/src/ott/math/_lbfgs.py b/src/ott/math/_lbfgs.py index 8a630d147..3620bb69e 100644 --- a/src/ott/math/_lbfgs.py +++ b/src/ott/math/_lbfgs.py @@ -11,10 +11,10 @@ # 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 typing import Any, Callable, Tuple +from collections.abc import Callable +from typing import Any import jax -import jax.numpy as jnp import optax @@ -25,11 +25,11 @@ def run_opt( opt: optax.GradientTransformationExtraArgs, - x_init: jnp.ndarray, - fun: Callable[[jnp.ndarray], jnp.ndarray], + x_init: jax.Array, + fun: Callable[[jax.Array], jax.Array], max_iter: int, tol: float, -) -> Tuple[jnp.ndarray, optax.OptState]: +) -> tuple[jax.Array, optax.OptState]: """Runs an optimization algorithm on a function. Args: @@ -68,12 +68,12 @@ def continuing_criterion(carry): def lbfgs( - fun: Callable[[jnp.ndarray], jnp.ndarray], - x_init: jnp.ndarray, + fun: Callable[[jax.Array], jax.Array], + x_init: jax.Array, max_iter: int = 100, tol: float = 1e-4, **kwargs: Any, -) -> Tuple[jnp.ndarray, optax.OptState]: +) -> tuple[jax.Array, optax.OptState]: """Runs optax's L-BFGS optimization on function. Args: diff --git a/src/ott/math/_legendre.py b/src/ott/math/_legendre.py index a6fe62b74..31fc5c529 100644 --- a/src/ott/math/_legendre.py +++ b/src/ott/math/_legendre.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Optional +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -22,9 +23,9 @@ def legendre( - fun: Callable[[jnp.ndarray], jnp.ndarray], + fun: Callable[[jax.Array], jax.Array], **kwargs: Any, -) -> Callable[[jnp.ndarray, Optional[jnp.ndarray], Any], jnp.ndarray]: +) -> Callable[[jax.Array, jax.Array | None, Any], jax.Array]: """Legendre (Fenchel) transform of a function. The solution is computed numerically using L-BFGS. @@ -42,8 +43,8 @@ def legendre( """ def fun_star( - x: jnp.ndarray, - x_init: Optional[jnp.ndarray] = None, + x: jax.Array, + x_init: jax.Array | None = None, ) -> float: """Runs optimization to compute the Legendre transform of ``fun`` at ``x``. @@ -57,7 +58,7 @@ def fun_star( """ x_init = x if x_init is None else x_init - def mod_fun(z: jnp.ndarray) -> float: + def mod_fun(z: jax.Array) -> float: """Conjugate maximizes - fun(z), here minimize fun(z) - .""" return fun(z) - jnp.dot(x, z) diff --git a/src/ott/math/_velocity_from_brenier_potential.py b/src/ott/math/_velocity_from_brenier_potential.py index 89d0e0f8e..e4840e8d0 100644 --- a/src/ott/math/_velocity_from_brenier_potential.py +++ b/src/ott/math/_velocity_from_brenier_potential.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -23,9 +24,9 @@ def velocity_from_brenier_potential( - potential: Callable[[jnp.ndarray], jnp.ndarray], + potential: Callable[[jax.Array], jax.Array], **kwargs: Any, -) -> Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]: +) -> Callable[[jax.Array, jax.Array], jax.Array]: """Get optimal time-dependent velocity field from :term:`Brenier potential`. The solution is computed numerically using a :term:`Legendre transform`. @@ -42,7 +43,7 @@ def velocity_from_brenier_potential( @functools.partial(jax.vmap, in_axes=[0, 0]) def vel(t: jnp.array, z: jnp.array) -> jnp.array: - def pot_t(x: jnp.ndarray) -> jnp.ndarray: + def pot_t(x: jax.Array) -> jax.Array: return 0.5 * (1 - t) * jnp.sum(x ** 2) + t * potential(x) grad_pot_t_star = jax.grad(math.legendre(pot_t, **kwargs)) diff --git a/src/ott/math/fixed_point_loop.py b/src/ott/math/fixed_point_loop.py index 764b20f85..da4b26aaf 100644 --- a/src/ott/math/fixed_point_loop.py +++ b/src/ott/math/fixed_point_loop.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -118,7 +119,7 @@ def fixpoint_iter_fwd( """ force_scan = min_iterations == max_iterations compute_error_flags = jnp.arange(inner_iterations) == inner_iterations - 1 - states = jax.tree_util.tree_map( + states = jax.tree.map( lambda x: jnp.zeros( (max_iterations // inner_iterations + 1,) + jnp.shape(x), dtype=jax.dtypes.result_type(x) @@ -136,7 +137,7 @@ def max_cond_fn(iteration_states_state): def unrolled_body_fn(iteration_states_state): iteration, states, state = iteration_states_state - states = jax.tree_util.tree_map( + states = jax.tree.map( lambda states, state: jax.lax.dynamic_update_index_in_dim( states, state, iteration // inner_iterations, 0 ), states, state @@ -177,9 +178,9 @@ def fixpoint_iter_bwd( force_scan = (min_iterations == max_iterations) constants, iteration, states = res # The tree may contain some python floats - g_constants = jax.tree_util.tree_map( + g_constants = jax.tree.map( lambda x: jnp.zeros_like(x, dtype=x.dtype) - if isinstance(x, (np.ndarray, jnp.ndarray)) else 0, constants + if isinstance(x, (np.ndarray, jax.Array)) else 0, constants ) def bwd_cond_fn(iteration_g_gconst): @@ -203,16 +204,12 @@ def one_iteration(iteration_state, compute_error): def unrolled_body_fn(iteration_g_gconst): iteration, g, g_constants = iteration_g_gconst - state = jax.tree_util.tree_map( - lambda x: x[iteration // inner_iterations], states - ) + state = jax.tree.map(lambda x: x[iteration // inner_iterations], states) _, pullback = jax.vjp( unrolled_body_fn_no_errors, iteration, constants, state ) _, gi_constants, g_state = pullback(g) - g_constants = jax.tree_util.tree_map( - lambda x, y: x + y, g_constants, gi_constants - ) + g_constants = jax.tree.map(lambda x, y: x + y, g_constants, gi_constants) out = (iteration - inner_iterations, g_state, g_constants) return (out, None) if force_scan else out diff --git a/src/ott/math/matrix_square_root.py b/src/ott/math/matrix_square_root.py index f017b6850..e97108bbd 100644 --- a/src/ott/math/matrix_square_root.py +++ b/src/ott/math/matrix_square_root.py @@ -13,7 +13,6 @@ # limitations under the License. import functools import math -from typing import Tuple import jax import jax.numpy as jnp @@ -32,13 +31,13 @@ ) ) def sqrtm( - x: jnp.ndarray, + x: jax.Array, threshold: float = 1e-6, min_iterations: int = 0, inner_iterations: int = 10, max_iterations: int = 1000, regularization: float = 1e-6 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array, jax.Array]: """Higham algorithm to compute matrix square root of p.d. matrix. See :cite:`higham:97`, eq. 2.6b @@ -123,10 +122,10 @@ def new_err(x, norm_x, y): def solve_sylvester_bartels_stewart( - a: jnp.ndarray, - b: jnp.ndarray, - c: jnp.ndarray, -) -> jnp.ndarray: + a: jax.Array, + b: jax.Array, + c: jax.Array, +) -> jax.Array: """Solve the real Sylvester equation AX - XB = C using Bartels-Stewart.""" # See https://nhigham.com/2020/09/01/what-is-the-sylvester-equation/ for # discussion of the algorithm (but note that in the derivation, the sign on @@ -159,14 +158,13 @@ def solve_sylvester_bartels_stewart( def sqrtm_fwd( - x: jnp.ndarray, + x: jax.Array, threshold: float, min_iterations: int, inner_iterations: int, max_iterations: int, regularization: float, -) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, - jnp.ndarray]]: +) -> tuple[tuple[jax.Array, jax.Array, jax.Array], tuple[jax.Array, jax.Array]]: """Forward pass of custom VJP.""" sqrt_x, inv_sqrt_x, errors = sqrtm( x=x, @@ -185,9 +183,9 @@ def sqrtm_bwd( inner_iterations: int, max_iterations: int, regularization: float, - residual: Tuple[jnp.ndarray, jnp.ndarray], - cotangent: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], -) -> Tuple[jnp.ndarray]: + residual: tuple[jax.Array, jax.Array], + cotangent: tuple[jax.Array, jax.Array, jax.Array], +) -> tuple[jax.Array]: """Compute the derivative by solving a Sylvester equation.""" del threshold, min_iterations, inner_iterations, \ max_iterations, regularization @@ -249,13 +247,13 @@ def sqrtm_bwd( ) ) def sqrtm_only( # noqa: D103 - x: jnp.ndarray, + x: jax.Array, threshold: float = 1e-6, min_iterations: int = 0, inner_iterations: int = 10, max_iterations: int = 1000, regularization: float = 1e-6 -) -> jnp.ndarray: +) -> jax.Array: return sqrtm( x, threshold, min_iterations, inner_iterations, max_iterations, regularization @@ -263,9 +261,9 @@ def sqrtm_only( # noqa: D103 def sqrtm_only_fwd( # noqa: D103 - x: jnp.ndarray, threshold: float, min_iterations: int, + x: jax.Array, threshold: float, min_iterations: int, inner_iterations: int, max_iterations: int, regularization: float -) -> Tuple[jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array]: sqrt_x = sqrtm( x, threshold, min_iterations, inner_iterations, max_iterations, regularization @@ -275,9 +273,9 @@ def sqrtm_only_fwd( # noqa: D103 def sqrtm_only_bwd( # noqa: D103 threshold: float, min_iterations: int, inner_iterations: int, - max_iterations: int, regularization: float, sqrt_x: jnp.ndarray, - cotangent: jnp.ndarray -) -> Tuple[jnp.ndarray]: + max_iterations: int, regularization: float, sqrt_x: jax.Array, + cotangent: jax.Array +) -> tuple[jax.Array]: del threshold, min_iterations, inner_iterations, \ max_iterations, regularization vjp = jnp.swapaxes( @@ -301,13 +299,13 @@ def sqrtm_only_bwd( # noqa: D103 ) ) def inv_sqrtm_only( # noqa: D103 - x: jnp.ndarray, + x: jax.Array, threshold: float = 1e-6, min_iterations: int = 0, inner_iterations: int = 10, max_iterations: int = 1000, regularization: float = 1e-6 -) -> jnp.ndarray: +) -> jax.Array: return sqrtm( x, threshold, min_iterations, inner_iterations, max_iterations, regularization @@ -315,13 +313,13 @@ def inv_sqrtm_only( # noqa: D103 def inv_sqrtm_only_fwd( # noqa: D103 - x: jnp.ndarray, + x: jax.Array, threshold: float, min_iterations: int, inner_iterations: int, max_iterations: int, regularization: float, -) -> Tuple[jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array]: inv_sqrt_x = sqrtm( x, threshold, min_iterations, inner_iterations, max_iterations, regularization @@ -331,9 +329,9 @@ def inv_sqrtm_only_fwd( # noqa: D103 def inv_sqrtm_only_bwd( # noqa: D103 threshold: float, min_iterations: int, inner_iterations: int, - max_iterations: int, regularization: float, residual: jnp.ndarray, - cotangent: jnp.ndarray -) -> Tuple[jnp.ndarray]: + max_iterations: int, regularization: float, residual: jax.Array, + cotangent: jax.Array +) -> tuple[jax.Array]: del threshold, min_iterations, inner_iterations, \ max_iterations, regularization diff --git a/src/ott/math/unbalanced_functions.py b/src/ott/math/unbalanced_functions.py index 22bf6800f..2bb824170 100644 --- a/src/ott/math/unbalanced_functions.py +++ b/src/ott/math/unbalanced_functions.py @@ -11,33 +11,34 @@ # 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 typing import Callable +from collections.abc import Callable +import jax import jax.numpy as jnp -def phi_star(h: jnp.ndarray, rho: float) -> jnp.ndarray: +def phi_star(h: jax.Array, rho: float) -> jax.Array: """Legendre transform of KL, :cite:`sejourne:19`, p. 9.""" return rho * (jnp.exp(h / rho) - 1) -def derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: +def derivative_phi_star(f: jax.Array, rho: float) -> jax.Array: """Derivative of Legendre transform of phi_starKL, see phi_star.""" # TODO(cuturi): use jax.grad directly. return jnp.exp(f / rho) def grad_of_marginal_fit( - c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float -) -> jnp.ndarray: + c: jax.Array, h: jax.Array, tau: float, epsilon: float +) -> jax.Array: """Compute grad of terms linked to marginals in objective. Computes gradient w.r.t. f ( or g) of terms in :cite:`sejourne:19`, left-hand-side of eq. 15 terms involving phi_star). Args: - c: jnp.ndarray, first target marginal (either a or b in practice) - h: jnp.ndarray, potential (either f or g in practice) + c: jax.Array, first target marginal (either a or b in practice) + h: jax.Array, potential (either f or g in practice) tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal epsilon: regularization @@ -50,14 +51,14 @@ def grad_of_marginal_fit( return jnp.where(c > 0, c * derivative_phi_star(-h, r), 0.0) -def second_derivative_phi_star(f: jnp.ndarray, rho: float) -> jnp.ndarray: +def second_derivative_phi_star(f: jax.Array, rho: float) -> jax.Array: """Second Derivative of Legendre transform of KL, see phi_star.""" return jnp.exp(f / rho) / rho def diag_jacobian_of_marginal_fit( - c: jnp.ndarray, h: jnp.ndarray, tau: float, epsilon: float, - derivative: Callable[[jnp.ndarray, float], jnp.ndarray] + c: jax.Array, h: jax.Array, tau: float, epsilon: float, + derivative: Callable[[jax.Array, float], jax.Array] ): """Compute grad of terms linked to marginals in objective. @@ -65,8 +66,8 @@ def diag_jacobian_of_marginal_fit( left-hand-side of eq. 32 (terms involving phi_star) Args: - c: jnp.ndarray, first target marginal (either a or b in practice) - h: jnp.ndarray, potential (either f or g in practice) + c: jax.Array, first target marginal (either a or b in practice) + h: jax.Array, potential (either f or g in practice) tau: float, strength (in ]0,1]) of regularizer w.r.t. marginal epsilon: regularization derivative: Callable diff --git a/src/ott/math/utils.py b/src/ott/math/utils.py index 3f30b969e..6305f3958 100644 --- a/src/ott/math/utils.py +++ b/src/ott/math/utils.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import TYPE_CHECKING import jax import jax.numpy as jnp @@ -37,10 +38,10 @@ def safe_log( # noqa: D103 - x: jnp.ndarray, + x: jax.Array, *, - eps: Optional[float] = None -) -> jnp.ndarray: + eps: float | None = None +) -> jax.Array: if eps is None: eps = jnp.finfo(x.dtype).tiny return jnp.where(x > 0.0, jnp.log(x), jnp.log(eps)) @@ -49,11 +50,11 @@ def safe_log( # noqa: D103 @functools.partial(jax.custom_jvp, nondiff_argnames=("ord", "axis", "keepdims")) @functools.partial(jax.jit, static_argnames=("ord", "axis", "keepdims")) def norm( - x: jnp.ndarray, - ord: Union[int, str, None] = None, - axis: Union[None, Sequence[int], int] = None, + x: jax.Array, + ord: int | str | None = None, + axis: None | Sequence[int] | int = None, keepdims: bool = False -) -> jnp.ndarray: +) -> jax.Array: """Computes order ord norm of vector, using `jnp.linalg` in forward pass. Evaluations of distances between a vector and itself using translation @@ -108,23 +109,23 @@ def norm_jvp(ord, axis, keepdims, primals, tangents): # TODO(michalk8): add axis argument -def kl(p: jnp.ndarray, q: jnp.ndarray) -> float: +def kl(p: jax.Array, q: jax.Array) -> float: """Kullback-Leibler divergence.""" return jnp.vdot(p, (safe_log(p) - safe_log(q))) -def gen_ent(x: jnp.ndarray) -> float: +def gen_ent(x: jax.Array) -> float: """Generalized entropy, adds the sum of ``x`` compared to usual entropy.""" return jnp.sum(jsp.special.entr(x)) + jnp.sum(x) -def gen_kl(p: jnp.ndarray, q: jnp.ndarray) -> float: +def gen_kl(p: jax.Array, q: jax.Array) -> float: """Generalized Kullback-Leibler divergence.""" return jnp.vdot(p, (safe_log(p) - safe_log(q))) + jnp.sum(q) - jnp.sum(p) # TODO(michalk8): add axis argument -def gen_js(p: jnp.ndarray, q: jnp.ndarray, c: float = 0.5) -> float: +def gen_js(p: jax.Array, q: jax.Array, c: float = 0.5) -> float: """Jensen-Shannon divergence.""" return c * (gen_kl(p, q) + gen_kl(q, p)) @@ -190,10 +191,10 @@ def logsumexp_jvp(axis, keepdims, return_sign, primals, tangents): @functools.partial(jax.custom_vjp, nondiff_argnames=("axis",)) def softmin( - x: jnp.ndarray, + x: jax.Array, gamma: float, - axis: Optional[Union[int, Sequence[int]]] = None -) -> jnp.ndarray: + axis: int | Sequence[int] | None = None +) -> jax.Array: r"""Soft-min operator. Args: @@ -221,8 +222,8 @@ def softmin( @functools.partial(jax.vmap, in_axes=[0, 0, None]) def barycentric_projection( - matrix: jnp.ndarray, y: jnp.ndarray, cost_fn: "costs.CostFn" -) -> jnp.ndarray: + matrix: jax.Array, y: jax.Array, cost_fn: "costs.CostFn" +) -> jax.Array: """Compute the barycentric projection of a matrix. Args: @@ -242,7 +243,7 @@ def sort_and_argsort( x: jnp.array, *, argsort: bool = False -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: +) -> tuple[jax.Array, jax.Array | None]: """Unified function that returns both sort and argsort, if latter needed.""" if argsort: i_x = jnp.argsort(x) @@ -251,9 +252,7 @@ def sort_and_argsort( @functools.partial(jax.custom_jvp, nondiff_argnames=("tol", "max_iter")) -def lambertw( - z: jnp.ndarray, tol: float = 1e-8, max_iter: int = 100 -) -> jnp.ndarray: +def lambertw(z: jax.Array, tol: float = 1e-8, max_iter: int = 100) -> jax.Array: """Principal branch of the `Lambert W function `_. @@ -269,7 +268,7 @@ def lambertw( The Lambert W evaluated at ``z``. """ # noqa: D205 - def initial_iacono(x: jnp.ndarray) -> jnp.ndarray: + def initial_iacono(x: jax.Array) -> jax.Array: y = jnp.sqrt(1.0 + jnp.e * x) num = 1.0 + 1.14956131 * y denom = 1.0 + 0.45495740 * jnp.log1p(y) @@ -302,9 +301,9 @@ def halley_iteration(container): @lambertw.defjvp def _lambertw_jvp( - tol: float, max_iter: int, primals: Tuple[jnp.ndarray, ...], - tangents: Tuple[jnp.ndarray, ...] -) -> Tuple[jnp.ndarray, jnp.ndarray]: + tol: float, max_iter: int, primals: tuple[jax.Array, ...], + tangents: tuple[jax.Array, ...] +) -> tuple[jax.Array, jax.Array]: z, = primals dz, = tangents w = lambertw(z, tol=tol, max_iter=max_iter) diff --git a/src/ott/neural/data/ot_dataloader.py b/src/ott/neural/data/ot_dataloader.py index 9fe4a02b2..3e8df7d04 100644 --- a/src/ott/neural/data/ot_dataloader.py +++ b/src/ott/neural/data/ot_dataloader.py @@ -13,7 +13,8 @@ # limitations under the License. import dataclasses import functools -from typing import Any, Iterable, Iterator, Literal, Optional, Tuple +from collections.abc import Iterable, Iterator +from typing import Any, Literal import jax import jax.random as jr @@ -49,14 +50,14 @@ class LinearOTDataloader: shardings: Input and output shardings for the source and target arrays. """ rng: jax.Array - dataset: Iterable[Tuple[jax.Array, jax.Array]] - epsilon: Optional[float] = None - relative_epsilon: Optional[Literal["mean", "std"]] = None - cost_fn: Optional[costs.CostFn] = None + dataset: Iterable[tuple[jax.Array, jax.Array]] + epsilon: float | None = None + relative_epsilon: Literal["mean", "std"] | None = None + cost_fn: costs.CostFn | None = None threshold: float = 1e-3 max_iterations: int = 2000 replace: bool = True - shardings: Optional[jax.sharding.Sharding] = None + shardings: jax.sharding.Sharding | None = None def __post_init__(self) -> None: self._align_fn = jax.jit( @@ -69,8 +70,8 @@ def __post_init__(self) -> None: in_shardings=(None, self.shardings, self.shardings), out_shardings=(self.shardings, self.shardings), ) - self._data_it: Optional[Iterator[Tuple[jax.Array, jax.Array]]] = None - self._rng_it: Optional[jax.Array] = None + self._data_it: Iterator[tuple[jax.Array, jax.Array]] | None = None + self._rng_it: jax.Array | None = None def __iter__(self) -> "LinearOTDataloader": """Return self.""" @@ -78,7 +79,7 @@ def __iter__(self) -> "LinearOTDataloader": self._rng_it = self.rng return self - def __next__(self) -> Tuple[jax.Array, jax.Array]: + def __next__(self) -> tuple[jax.Array, jax.Array]: """Align source and target samples in a batch. Returns: @@ -104,11 +105,11 @@ def _align( x: jax.Array, y: jax.Array, cost_fn: costs.CostFn, - epsilon: Optional[float], - relative_epsilon: Optional[Literal["mean", "std"]], + epsilon: float | None, + relative_epsilon: Literal["mean", "std"] | None, replace: bool, **kwargs: Any, -) -> Tuple[jax.Array, jax.Array]: +) -> tuple[jax.Array, jax.Array]: geom = pointcloud.PointCloud( x, y, diff --git a/src/ott/neural/data/semidiscrete_dataloader.py b/src/ott/neural/data/semidiscrete_dataloader.py index 08ca26bfc..ca5ccb2d6 100644 --- a/src/ott/neural/data/semidiscrete_dataloader.py +++ b/src/ott/neural/data/semidiscrete_dataloader.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Optional, Tuple, Union import jax import jax.numpy as jnp @@ -57,11 +56,11 @@ class SemidiscreteDataloader: rng: jax.Array sd_out: semidiscrete.SemidiscreteOutput batch_size: int - epsilon: Optional[float] = None - subset_size_threshold: Optional[int] = None - subset_size: Optional[int] = None + epsilon: float | None = None + subset_size_threshold: int | None = None + subset_size: int | None = None return_indices: bool = False - out_shardings: Optional[jax.sharding.Sharding] = None + out_shardings: jax.sharding.Sharding | None = None def __post_init__(self) -> None: _, m = self.sd_out.geom.shape @@ -75,7 +74,7 @@ def __post_init__(self) -> None: assert 0 < self.subset_size < m, \ f"Subset size must be in (0, {m}), got {self.subset_size}." - self._rng_it: Optional[jax.Array] = None + self._rng_it: jax.Array | None = None self._sample_fn = jax.jit( _sample, out_shardings=self.out_shardings, @@ -95,8 +94,7 @@ def __iter__(self) -> "SemidiscreteDataloader": def __next__( self - ) -> Union[Tuple[jax.Array, jax.Array], Tuple[jax.Array, jax.Array, - jax.Array]]: + ) -> tuple[jax.Array, jax.Array] | tuple[jax.Array, jax.Array, jax.Array]: """Sample from the source distribution and match it with the data. Returns: @@ -120,11 +118,11 @@ def _sample( rng: jax.Array, out: semidiscrete.SemidiscreteOutput, batch_size: int, - epsilon: Optional[float], - subset_size_threshold: Optional[int], + epsilon: float | None, + subset_size_threshold: int | None, subset_size: int, return_indices: bool, -) -> Union[Tuple[jax.Array, jax.Array], Tuple[jax.Array, jax.Array, jax.Array]]: +) -> tuple[jax.Array, jax.Array] | tuple[jax.Array, jax.Array, jax.Array]: rng_sample, rng_tmat = jr.split(rng, 2) out_sampled = out.sample(rng_sample, batch_size, epsilon=epsilon) @@ -148,7 +146,7 @@ def _sample_from_coupling( rng: jax.Array, coupling: jax.Array, *, - subset_size_threshold: Optional[int], + subset_size_threshold: int | None, subset_size: int, axis: int, ) -> jax.Array: diff --git a/src/ott/neural/methods/conditional_monge_gap.py b/src/ott/neural/methods/conditional_monge_gap.py index 9ff80db5d..c85457c18 100644 --- a/src/ott/neural/methods/conditional_monge_gap.py +++ b/src/ott/neural/methods/conditional_monge_gap.py @@ -14,17 +14,8 @@ import collections import functools import logging -from typing import ( - Any, - Callable, - Dict, - Iterator, - Literal, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Iterator, Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -50,18 +41,18 @@ def cmonge_gap_from_samples( - source: jnp.ndarray, - target: jnp.ndarray, - condition: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_cost", "median"]] = 1.0, + source: jax.Array, + target: jax.Array, + condition: jax.Array, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, + relative_epsilon: Literal["mean", "std"] | None = None, + scale_cost: float | Literal["mean", "max_cost", "median"] = 1.0, return_output: bool = False, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, + num_segments: int | None = None, + max_measure_size: int | None = None, **kwargs: Any, -) -> Union[float, Tuple[float, jnp.ndarray]]: +) -> float | tuple[float, jax.Array]: r"""Conditional Monge gap from samples using the segment interface. Computes the average Monge gap across conditions: @@ -133,11 +124,11 @@ def cmonge_gap_from_samples( # ott.neural.methods.monge_gap.monge_gap_from_samples` # as well as `ott.geometry.segment.py` def eval_fn( - padded_x: jnp.ndarray, - padded_y: jnp.ndarray, - padded_weight_x: jnp.ndarray, - padded_weight_y: jnp.ndarray, - ) -> jnp.ndarray: + padded_x: jax.Array, + padded_y: jax.Array, + padded_weight_x: jax.Array, + padded_weight_y: jax.Array, + ) -> jax.Array: """Monge gap for a single (padded) condition segment.""" # Displacement cost: weighted mean of pairwise costs c(x_i, T(x_i)). # Padded entries have weight 0, so they do not contribute. @@ -218,18 +209,16 @@ def __init__( self, dim_data: int, model: ConditionalPerturbationNetwork, - optimizer: Optional[optax.OptState] = None, - fitting_loss: Optional[Callable[[jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]]]] = None, - regularizer: Optional[Callable[ - [jnp.ndarray, jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]], - ]] = None, - regularizer_strength: Union[float, Sequence[float]] = 1.0, + optimizer: optax.OptState | None = None, + fitting_loss: Callable[[jax.Array, jax.Array], tuple[float, Any | None]] + | None = None, + regularizer: Callable[[jax.Array, jax.Array, jax.Array], + tuple[float, Any | None]] | None = None, + regularizer_strength: float | Sequence[float] = 1.0, num_train_iters: int = 10_000, logging: bool = False, valid_freq: int = 500, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, ): self._fitting_loss = fitting_loss self._regularizer = regularizer @@ -264,8 +253,8 @@ def setup( @property def regularizer( self, - ) -> Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray], Tuple[float, - Optional[Any]]]: + ) -> Callable[[jax.Array, jax.Array, jax.Array], tuple[float, Any + | None]]: """Conditional regularizer ``(source, mapped, labels) -> (loss, log)``. Defaults to zero if not provided. @@ -277,7 +266,7 @@ def regularizer( @property def fitting_loss( self, - ) -> Callable[[jnp.ndarray, jnp.ndarray], Tuple[float, Optional[Any]]]: + ) -> Callable[[jax.Array, jax.Array], tuple[float, Any | None]]: """Fitting loss ``(mapped, target) -> (loss, log)``. Defaults to zero if not provided. @@ -288,11 +277,11 @@ def fitting_loss( @staticmethod def _generate_batch( - loader_source: Iterator[jnp.ndarray], - loader_target: Iterator[jnp.ndarray], - loader_condition: Iterator[jnp.ndarray], - loader_label: Iterator[jnp.ndarray], - ) -> Dict[str, jnp.ndarray]: + loader_source: Iterator[jax.Array], + loader_target: Iterator[jax.Array], + loader_condition: Iterator[jax.Array], + loader_label: Iterator[jax.Array], + ) -> dict[str, jax.Array]: """Generate a batch of samples from all four iterators.""" return { "source": next(loader_source), @@ -303,15 +292,15 @@ def _generate_batch( def train_map_estimator( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - trainloader_condition: Iterator[jnp.ndarray], - trainloader_label: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - validloader_condition: Iterator[jnp.ndarray], - validloader_label: Iterator[jnp.ndarray], - ) -> Tuple[train_state.TrainState, Dict[str, Any]]: + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + trainloader_condition: Iterator[jax.Array], + trainloader_label: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + validloader_condition: Iterator[jax.Array], + validloader_label: Iterator[jax.Array], + ) -> tuple[train_state.TrainState, dict[str, Any]]: """Training loop.""" logs = collections.defaultdict(lambda: collections.defaultdict(list)) @@ -372,9 +361,9 @@ def _get_step_fn(self) -> Callable: def loss_fn( params: frozen_dict.FrozenDict, apply_fn: Callable, - batch: Dict[str, jnp.ndarray], + batch: dict[str, jax.Array], step: int, - ) -> Tuple[float, Dict[str, float]]: + ) -> tuple[float, dict[str, float]]: """Loss function with conditional map and regularizer.""" # Apply the conditional map: T(source, condition) mapped_samples = apply_fn({"params": params}, batch["source"], @@ -407,11 +396,11 @@ def loss_fn( @functools.partial(jax.jit, static_argnums=3) def step_fn( state_neural_net: train_state.TrainState, - train_batch: Dict[str, jnp.ndarray], - valid_batch: Optional[Dict[str, jnp.ndarray]] = None, + train_batch: dict[str, jax.Array], + valid_batch: dict[str, jax.Array] | None = None, is_logging_step: bool = False, step: int = 0, - ) -> Tuple[train_state.TrainState, Dict[str, float]]: + ) -> tuple[train_state.TrainState, dict[str, float]]: """One step function.""" grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) (_, current_train_logs), grads = grad_fn( diff --git a/src/ott/neural/methods/expectile_neural_dual.py b/src/ott/neural/methods/expectile_neural_dual.py index d8ebbbff3..cf8e0ceef 100644 --- a/src/ott/neural/methods/expectile_neural_dual.py +++ b/src/ott/neural/methods/expectile_neural_dual.py @@ -11,16 +11,8 @@ # 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 typing import ( - Callable, - Dict, - Iterator, - List, - Literal, - Optional, - Tuple, - Union, -) +from collections.abc import Callable, Iterator +from typing import Literal import jax import jax.numpy as jnp @@ -37,7 +29,7 @@ __all__ = ["ENOTPotentials", "PotentialModelWrapper", "ExpectileNeuralDual"] -Train_t = Dict[Literal["train_logs", "valid_logs"], Dict[str, List[float]]] +Train_t = dict[Literal["train_logs", "valid_logs"], dict[str, list[float]]] Callback_t = Callable[[int, dual_potentials.DualPotentials], None] @@ -60,7 +52,7 @@ def __init__( ): self.__grad_f = grad_f - def f_potential(x: jnp.ndarray) -> jnp.ndarray: + def f_potential(x: jax.Array) -> jax.Array: y_hat = cost_fn.twist_operator(x, grad_f(x), False) y_hat = jax.lax.stop_gradient(y_hat) return -g(y_hat) + cost_fn(x, y_hat) @@ -68,7 +60,7 @@ def f_potential(x: jnp.ndarray) -> jnp.ndarray: super().__init__(f_potential, g, cost_fn=cost_fn) @property - def _grad_f(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + def _grad_f(self) -> Callable[[jax.Array], jax.Array]: return jax.vmap(self.__grad_f) @@ -87,15 +79,15 @@ class PotentialModelWrapper(potentials.BasePotential): is_potential: bool = True @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + def __call__(self, x: jax.Array) -> jax.Array: """Apply model and optionally add l2 norm or x.""" - z: jnp.ndarray = self.model(x) + z: jax.Array = self.model(x) if self.is_potential: z = z.squeeze() return z def potential_gradient_fn( - self, params: frozen_dict.FrozenDict[str, jnp.ndarray] + self, params: frozen_dict.FrozenDict[str, jax.Array] ) -> potentials.PotentialGradientFn_t: """A vector function or gradient of the potential.""" if self.is_potential: @@ -156,18 +148,18 @@ class ExpectileNeuralDual: def __init__( self, dim_data: int, - neural_f: Optional[nn.Module] = None, - neural_g: Optional[nn.Module] = None, - optimizer_f: Optional[optax.GradientTransformation] = None, - optimizer_g: Optional[optax.GradientTransformation] = None, - cost_fn: Optional[costs.TICost] = None, + neural_f: nn.Module | None = None, + neural_g: nn.Module | None = None, + optimizer_f: optax.GradientTransformation | None = None, + optimizer_g: optax.GradientTransformation | None = None, + cost_fn: costs.TICost | None = None, expectile: float = 0.99, expectile_loss_coef: float = 1.0, num_train_iters: int = 20000, valid_freq: int = 1000, log_freq: int = 1000, logging: bool = False, - rng: Optional[jax.Array] = None + rng: jax.Array | None = None ): self.num_train_iters = num_train_iters self.valid_freq = valid_freq @@ -212,12 +204,12 @@ def __init__( def __call__( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, - ) -> Union[ENOTPotentials, Tuple[ENOTPotentials, Train_t]]: + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + callback: Callback_t | None = None, + ) -> ENOTPotentials | tuple[ENOTPotentials, Train_t]: """Train and return the Kantorovich dual potentials.""" logs = self.train_fn( trainloader_source, @@ -232,11 +224,11 @@ def __call__( def train_fn( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + callback: Callback_t | None = None, ) -> Train_t: """Training and validation.""" try: @@ -288,10 +280,10 @@ def train_fn( def _get_train_step( self ) -> Callable[[ - potentials.PotentialTrainState, potentials.PotentialTrainState, Dict[ - str, jnp.ndarray] - ], Tuple[potentials.PotentialTrainState, potentials.PotentialTrainState, - jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]]: + potentials.PotentialTrainState, potentials.PotentialTrainState, dict[ + str, jax.Array] + ], tuple[potentials.PotentialTrainState, potentials.PotentialTrainState, + jax.Array, jax.Array, jax.Array, jax.Array]]: @jax.jit def step_fn(state_f, state_g, batch): @@ -314,9 +306,9 @@ def step_fn(state_f, state_g, batch): def _get_valid_step( self ) -> Callable[[ - potentials.PotentialTrainState, potentials.PotentialTrainState, Dict[ - str, jnp.ndarray] - ], Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]]: + potentials.PotentialTrainState, potentials.PotentialTrainState, dict[ + str, jax.Array] + ], tuple[jax.Array, jax.Array, jax.Array]]: @jax.jit def step_fn(state_f, state_g, batch): @@ -332,14 +324,14 @@ def step_fn(state_f, state_g, batch): return step_fn - def _expectile_loss(self, diff: jnp.ndarray) -> jnp.ndarray: + def _expectile_loss(self, diff: jax.Array) -> jax.Array: """Loss of the expectile regression :cite:`buzun:24`.""" weight = jnp.where(diff >= 0, self.expectile, (1 - self.expectile)) return weight * diff ** 2 def _get_g_value_partial( - self, params_g: frozen_dict.FrozenDict[str, jnp.ndarray], - g_value: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], + self, params_g: frozen_dict.FrozenDict[str, jax.Array], + g_value: Callable[[frozen_dict.FrozenDict[str, jax.Array]], potentials.PotentialValueFn_t] ): @@ -349,14 +341,14 @@ def _get_g_value_partial( return g_value_partial, g_value_partial_detach def _loss_fn( - self, params_f: frozen_dict.FrozenDict[str, jnp.ndarray], - params_g: frozen_dict.FrozenDict[str, jnp.ndarray], - gradient_f: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], + self, params_f: frozen_dict.FrozenDict[str, jax.Array], + params_g: frozen_dict.FrozenDict[str, jax.Array], + gradient_f: Callable[[frozen_dict.FrozenDict[str, jax.Array]], potentials.PotentialGradientFn_t], - g_value: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], - potentials.PotentialValueFn_t], batch: Dict[str, - jnp.ndarray] - ) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]]: + g_value: Callable[[frozen_dict.FrozenDict[str, jax.Array]], + potentials.PotentialValueFn_t], batch: dict[str, + jax.Array] + ) -> tuple[jax.Array, tuple[jax.Array, jax.Array, jax.Array]]: source, target = batch["source"], batch["target"] @@ -414,10 +406,10 @@ def to_dual_potentials(self) -> ENOTPotentials: @staticmethod def _update_logs( - logs: Dict[str, List[Union[float, str]]], - loss_f: jnp.ndarray, - loss_g: jnp.ndarray, - w_dist: jnp.ndarray, + logs: dict[str, list[float | str]], + loss_f: jax.Array, + loss_g: jax.Array, + w_dist: jax.Array, ) -> None: logs["loss_f"].append(float(loss_f)) logs["loss_g"].append(float(loss_g)) diff --git a/src/ott/neural/methods/flow_matching.py b/src/ott/neural/methods/flow_matching.py index 57b1ee777..30dc7d2c6 100644 --- a/src/ott/neural/methods/flow_matching.py +++ b/src/ott/neural/methods/flow_matching.py @@ -13,22 +13,15 @@ # limitations under the License. import functools import inspect -from typing import ( - Any, - Callable, - Dict, - Literal, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Sequence +from typing import Any, Literal import jax import jax.numpy as jnp import jax.random as jr import jax.tree_util as jtu import numpy as np +from jax.typing import DTypeLike import diffrax import optax @@ -42,8 +35,8 @@ "gaussian_nll", ] -DivState = Tuple[jax.Array, jax.Array] # velocity, divergence -Batch = Dict[Literal["t", "x_t", "v_t", "cond"], jax.Array] +DivState = tuple[jax.Array, jax.Array] # velocity, divergence +Batch = dict[Literal["t", "x_t", "v_t", "cond"], jax.Array] def flow_matching_step( @@ -52,9 +45,9 @@ def flow_matching_step( batch: Batch, *, loss_fn: Callable[[jax.Array, jax.Array], jax.Array] = optax.squared_error, - model_callback_fn: Optional[Callable[[nnx.Module], None]] = None, - rngs: Optional[nnx.Rngs] = None, -) -> Dict[Literal["loss", "grad_norm"], jax.Array]: + model_callback_fn: Callable[[nnx.Module], None] | None = None, + rngs: nnx.Rngs | None = None, +) -> dict[Literal["loss", "grad_norm"], jax.Array]: """Perform a flow matching step. Args: @@ -99,10 +92,10 @@ def interpolate_samples( rng: jax.Array, x0: jax.Array, x1: jax.Array, - cond: Optional[jax.Array] = None, + cond: jax.Array | None = None, *, - time_sampler: Optional[Callable[[jax.Array, Tuple[int], jnp.dtype], - jax.Array]] = None + time_sampler: Callable[[jax.Array, tuple[int], DTypeLike], jax.Array] + | None = None ) -> Batch: """Sample time and interpolate. @@ -142,16 +135,16 @@ def interpolate_samples( def evaluate_velocity_field( model: nnx.Module, - x: Union[jax.Array, Any], - cond: Optional[jax.Array] = None, + x: jax.Array | Any, + cond: jax.Array | None = None, *, t0: float = 0.0, t1: float = 1.0, reverse: bool = False, - num_steps: Optional[int] = None, - solver: Optional[diffrax.AbstractSolver] = None, - save_trajectory_kwargs: Optional[Dict[str, Any]] = None, - save_velocity_kwargs: Optional[Dict[str, Any]] = None, + num_steps: int | None = None, + solver: diffrax.AbstractSolver | None = None, + save_trajectory_kwargs: dict[str, Any] | None = None, + save_velocity_kwargs: dict[str, Any] | None = None, **kwargs: Any, ) -> diffrax.Solution: """Solve an ODE. @@ -226,13 +219,13 @@ def evaluate_velocity_field( def curvature( model: nnx.Module, x0: jax.Array, - cond: Optional[jax.Array] = None, + cond: jax.Array | None = None, *, - ts: Union[int, jax.Array, Sequence[float]], - drop_last_velocity: Optional[bool] = None, + ts: int | jax.Array | Sequence[float], + drop_last_velocity: bool | None = None, loss_fn: Callable[[jax.Array, jax.Array], jax.Array] = optax.squared_error, **kwargs: Any, -) -> Tuple[jax.Array, diffrax.Solution]: +) -> tuple[jax.Array, diffrax.Solution]: """Compute the curvature :cite:`lee:23`. Also known as straightness in :cite:`liu:22`. @@ -284,12 +277,12 @@ def curvature( def gaussian_nll( model: nnx.Module, x1: jax.Array, - cond: Optional[jax.Array] = None, + cond: jax.Array | None = None, *, - noise: Optional[jax.Array] = None, + noise: jax.Array | None = None, stddev: float = 1.0, **kwargs: Any, -) -> Tuple[jax.Array, diffrax.Solution]: +) -> tuple[jax.Array, diffrax.Solution]: """Compute the Gaussian negative log-likelihood. Args: @@ -336,19 +329,19 @@ def gaussian_nll( def _velocity( - t: jax.Array, x_t: jax.Array, cond: Optional[jax.Array], model: nnx.Module + t: jax.Array, x_t: jax.Array, cond: jax.Array | None, model: nnx.Module ) -> jax.Array: cond = None if cond is None else cond[None] return model(t[None], x_t[None], cond).squeeze(0) def _exact_divergence( - t: jax.Array, state_t: DivState, cond: Optional[jax.Array], *, + t: jax.Array, state_t: DivState, cond: jax.Array | None, *, model: nnx.Module ) -> DivState: def divergence_v( - t: jax.Array, x: jax.Array, cond: Optional[jax.Array] + t: jax.Array, x: jax.Array, cond: jax.Array | None ) -> jax.Array: # divergence of fwd velocity field jacobian = jax.jacrev(_velocity, argnums=1)(t, x, cond, model) @@ -362,7 +355,7 @@ def divergence_v( def _hutchinson_divergence( - t: jax.Array, state_t: DivState, cond: Optional[jax.Array], *, + t: jax.Array, state_t: DivState, cond: jax.Array | None, *, model: nnx.Module, h: jax.Array ) -> DivState: x_t, _ = state_t diff --git a/src/ott/neural/methods/monge_gap.py b/src/ott/neural/methods/monge_gap.py index 56441795e..42f989706 100644 --- a/src/ott/neural/methods/monge_gap.py +++ b/src/ott/neural/methods/monge_gap.py @@ -13,17 +13,8 @@ # limitations under the License. import collections import functools -from typing import ( - Any, - Callable, - Dict, - Iterator, - Literal, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Iterator, Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -42,15 +33,15 @@ def monge_gap( - map_fn: Callable[[jnp.ndarray], jnp.ndarray], - reference_points: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_cost", "median"]] = 1.0, + map_fn: Callable[[jax.Array], jax.Array], + reference_points: jax.Array, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, + relative_epsilon: Literal["mean", "std"] | None = None, + scale_cost: float | Literal["mean", "max_cost", "median"] = 1.0, return_output: bool = False, **kwargs: Any -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: +) -> float | tuple[float, sinkhorn.SinkhornOutput]: r"""Monge gap regularizer :cite:`uscidda:23`. For a cost function :math:`c` and empirical reference measure @@ -107,15 +98,15 @@ def monge_gap( def monge_gap_from_samples( - source: jnp.ndarray, - target: jnp.ndarray, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - scale_cost: Union[float, Literal["mean", "max_cost", "median"]] = 1.0, + source: jax.Array, + target: jax.Array, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, + relative_epsilon: Literal["mean", "std"] | None = None, + scale_cost: float | Literal["mean", "max_cost", "median"] = 1.0, return_output: bool = False, **kwargs: Any -) -> Union[float, Tuple[float, sinkhorn.SinkhornOutput]]: +) -> float | tuple[float, sinkhorn.SinkhornOutput]: r"""Monge gap, instantiated in terms of samples before / after applying map. .. math:: @@ -209,16 +200,16 @@ def __init__( self, dim_data: int, model: potentials.BasePotential, - optimizer: Optional[optax.OptState] = None, - fitting_loss: Optional[Callable[[jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]]]] = None, - regularizer: Optional[Callable[[jnp.ndarray, jnp.ndarray], - Tuple[float, Optional[Any]]]] = None, - regularizer_strength: Union[float, Sequence[float]] = 1.0, + optimizer: optax.OptState | None = None, + fitting_loss: Callable[[jax.Array, jax.Array], tuple[float, Any | None]] + | None = None, + regularizer: Callable[[jax.Array, jax.Array], tuple[float, Any | None]] + | None = None, + regularizer_strength: float | Sequence[float] = 1.0, num_train_iters: int = 10_000, logging: bool = False, valid_freq: int = 500, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, ): self._fitting_loss = fitting_loss self._regularizer = regularizer @@ -257,7 +248,7 @@ def setup( self.step_fn = self._get_step_fn() @property - def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + def regularizer(self) -> Callable[[jax.Array, jax.Array], float]: """Regularizer added to the fitting loss. Can be, e.g. the @@ -271,7 +262,7 @@ def regularizer(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: return lambda *_, **__: (0.0, None) @property - def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: + def fitting_loss(self) -> Callable[[jax.Array, jax.Array], float]: """Fitting loss to fit the marginal constraint. Can be, e.g. :func:`~ott.tools.sinkhorn_divergence.sinkdiv`. @@ -284,9 +275,9 @@ def fitting_loss(self) -> Callable[[jnp.ndarray, jnp.ndarray], float]: @staticmethod def _generate_batch( - loader_source: Iterator[jnp.ndarray], - loader_target: Iterator[jnp.ndarray], - ) -> Dict[str, jnp.ndarray]: + loader_source: Iterator[jax.Array], + loader_target: Iterator[jax.Array], + ) -> dict[str, jax.Array]: """Generate batches a batch of samples. ``loader_source`` and ``loader_target`` can be training or @@ -299,11 +290,11 @@ def _generate_batch( def train_map_estimator( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - ) -> Tuple[train_state.TrainState, Dict[str, Any]]: + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + ) -> tuple[train_state.TrainState, dict[str, Any]]: """Training loop.""" # define logs logs = collections.defaultdict(lambda: collections.defaultdict(list)) @@ -361,8 +352,8 @@ def _get_step_fn(self) -> Callable: def loss_fn( params: frozen_dict.FrozenDict, apply_fn: Callable, - batch: Dict[str, jnp.ndarray], step: int - ) -> Tuple[float, Dict[str, float]]: + batch: dict[str, jax.Array], step: int + ) -> tuple[float, dict[str, float]]: """Loss function.""" # map samples with the fitted map mapped_samples = apply_fn({"params": params}, batch["source"]) @@ -392,11 +383,11 @@ def loss_fn( @functools.partial(jax.jit, static_argnums=3) def step_fn( state_neural_net: train_state.TrainState, - train_batch: Dict[str, jnp.ndarray], - valid_batch: Optional[Dict[str, jnp.ndarray]] = None, + train_batch: dict[str, jax.Array], + valid_batch: dict[str, jax.Array] | None = None, is_logging_step: bool = False, step: int = 0 - ) -> Tuple[train_state.TrainState, Dict[str, float]]: + ) -> tuple[train_state.TrainState, dict[str, float]]: """One step function.""" # compute loss and gradients grad_fn = jax.value_and_grad(loss_fn, argnums=0, has_aux=True) diff --git a/src/ott/neural/methods/neuraldual.py b/src/ott/neural/methods/neuraldual.py index caa63db75..6ebef76b4 100644 --- a/src/ott/neural/methods/neuraldual.py +++ b/src/ott/neural/methods/neuraldual.py @@ -12,16 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import warnings -from typing import ( - Callable, - Dict, - Iterator, - List, - Literal, - Optional, - Tuple, - Union, -) +from collections.abc import Callable, Iterator +from typing import Literal import jax import jax.numpy as jnp @@ -37,7 +29,7 @@ __all__ = ["W2NeuralDual"] -Train_t = Dict[Literal["train_logs", "valid_logs"], Dict[str, List[float]]] +Train_t = dict[Literal["train_logs", "valid_logs"], dict[str, list[float]]] Callback_t = Callable[[int, dual_potentials.DualPotentials], None] PotentialValueFn_t = potentials.PotentialValueFn_t @@ -46,7 +38,7 @@ def _value_fn( model: nnx.Module, - other_value_fn: Optional[Callable] = None, + other_value_fn: Callable | None = None, ) -> PotentialValueFn_t: """Get a scalar value function from an NNX model. @@ -60,7 +52,7 @@ def _value_fn( "The value of a gradient-based potential depends on the other potential." ) - def value_fn(x: jnp.ndarray) -> jnp.ndarray: + def value_fn(x: jax.Array) -> jax.Array: squeeze = x.ndim == 1 if squeeze: x = jnp.expand_dims(x, 0) @@ -147,19 +139,19 @@ class W2NeuralDual: def __init__( self, dim_data: int, - neural_f: Optional[nnx.Module] = None, - neural_g: Optional[nnx.Module] = None, - optimizer_f: Optional[optax.OptState] = None, - optimizer_g: Optional[optax.OptState] = None, + neural_f: nnx.Module | None = None, + neural_g: nnx.Module | None = None, + optimizer_f: optax.OptState | None = None, + optimizer_g: optax.OptState | None = None, num_train_iters: int = 20000, num_inner_iters: int = 1, - back_and_forth: Optional[bool] = None, + back_and_forth: bool | None = None, valid_freq: int = 1000, log_freq: int = 1000, logging: bool = False, - rng: Optional[jax.Array] = None, - conjugate_solver: Optional[conjugate.FenchelConjugateSolver - ] = conjugate.DEFAULT_CONJUGATE_SOLVER, + rng: jax.Array | None = None, + conjugate_solver: conjugate.FenchelConjugateSolver + | None = conjugate.DEFAULT_CONJUGATE_SOLVER, amortization_loss: Literal["objective", "regression"] = "regression", parallel_updates: bool = True, ): @@ -246,13 +238,13 @@ def setup( def __call__( # noqa: D102 self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, - ) -> Union[dual_potentials.DualPotentials, - Tuple[dual_potentials.DualPotentials, Train_t]]: + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + callback: Callback_t | None = None, + ) -> (dual_potentials.DualPotentials | + tuple[dual_potentials.DualPotentials, Train_t]): logs = self.train_fn( trainloader_source, trainloader_target, @@ -270,8 +262,8 @@ def _compute_losses( self, model_f: nnx.Module, model_g: nnx.Module, - batch: Dict[str, jnp.ndarray], - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + batch: dict[str, jax.Array], + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Compute all losses. Returns: @@ -282,7 +274,7 @@ def _compute_losses( g_gradient = _gradient_fn(model_g) init_source_hat = g_gradient(target) - def g_value_partial(y: jnp.ndarray) -> jnp.ndarray: + def g_value_partial(y: jax.Array) -> jax.Array: return _value_fn(model_g)(y) f_value_partial = _value_fn(model_f, g_value_partial) @@ -417,11 +409,11 @@ def valid_step(model_f, model_g, batch): def train_neuraldual_parallel( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + callback: Callback_t | None = None, ) -> Train_t: """Training and validation with parallel updates.""" try: @@ -488,11 +480,11 @@ def train_neuraldual_parallel( def train_neuraldual_alternating( self, - trainloader_source: Iterator[jnp.ndarray], - trainloader_target: Iterator[jnp.ndarray], - validloader_source: Iterator[jnp.ndarray], - validloader_target: Iterator[jnp.ndarray], - callback: Optional[Callback_t] = None, + trainloader_source: Iterator[jax.Array], + trainloader_target: Iterator[jax.Array], + validloader_source: Iterator[jax.Array], + validloader_target: Iterator[jax.Array], + callback: Callback_t | None = None, ) -> Train_t: """Training and validation with alternating updates.""" try: @@ -565,7 +557,7 @@ def to_dual_potentials( f_value = _value_fn(self.neural_f) g_value_prediction = _value_fn(self.neural_g, f_value) - def g_value_finetuned(y: jnp.ndarray) -> jnp.ndarray: + def g_value_finetuned(y: jax.Array) -> jax.Array: x_hat = jax.grad(g_value_prediction)(y) grad_g_y = jax.lax.stop_gradient( self.conjugate_solver.solve(f_value, y, x_init=x_hat).grad @@ -586,10 +578,10 @@ def g_value_finetuned(y: jnp.ndarray) -> jnp.ndarray: @staticmethod def _update_logs( - logs: Dict[str, List[Union[float, str]]], - loss_f: jnp.ndarray, - loss_g: jnp.ndarray, - w_dist: jnp.ndarray, + logs: dict[str, list[float | str]], + loss_f: jax.Array, + loss_g: jax.Array, + w_dist: jax.Array, ) -> None: logs["loss_f"].append(float(loss_f)) logs["loss_g"].append(float(loss_g)) diff --git a/src/ott/neural/networks/conditional_perturbation_network.py b/src/ott/neural/networks/conditional_perturbation_network.py index bafb73260..0da45e0ff 100644 --- a/src/ott/neural/networks/conditional_perturbation_network.py +++ b/src/ott/neural/networks/conditional_perturbation_network.py @@ -1,14 +1,7 @@ -from typing import ( - Any, - Callable, - Dict, - Iterable, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Iterable, Sequence +from typing import Any +import jax import jax.numpy as jnp import flax.linen as nn @@ -26,13 +19,13 @@ class ConditionalPerturbationNetwork(BasePotential): # Same length as context_entity_bonds if embed_cond_equal is False # (if True, first item is size of deep set layer, rest is ignored) dim_cond_map: Iterable[int] = (50,) - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.gelu + act_fn: Callable[[jax.Array], jax.Array] = nn.gelu is_potential: bool = False layer_norm: bool = False embed_cond_equal: bool = ( False # Whether all context variables should be treated as set or not ) - context_entity_bonds: Iterable[Tuple[int, int]] = ( + context_entity_bonds: Iterable[tuple[int, int]] = ( (0, 10), (10, 20), ) # (start, stop) slicing bounds per context modality in c; @@ -42,9 +35,9 @@ class ConditionalPerturbationNetwork(BasePotential): @nn.compact def __call__( self, - x: jnp.ndarray, - c: Optional[jnp.ndarray] = None - ) -> Union[jnp.ndarray, Dict[str, jnp.ndarray]]: # noqa: D102 + x: jax.Array, + c: jax.Array | None = None + ) -> jax.Array | dict[str, jax.Array]: # noqa: D102 """Forward pass: map (x, c) -> x + residual. Args: @@ -123,7 +116,7 @@ def __call__( def create_train_state( self, - rng: jnp.ndarray, + rng: jax.Array, optimizer: optax.OptState, dim_data: int, **kwargs: Any, diff --git a/src/ott/neural/networks/icnn.py b/src/ott/neural/networks/icnn.py index a28bcccdf..12b4261d0 100644 --- a/src/ott/neural/networks/icnn.py +++ b/src/ott/neural/networks/icnn.py @@ -13,7 +13,7 @@ # limitations under the License. """Input convex neural networks and KeyNet (vector-output variant).""" -from typing import Callable, Optional, Sequence, Tuple, Union +from collections.abc import Callable, Sequence import jax import jax.numpy as jnp @@ -39,9 +39,9 @@ def _get_act_alpha(act_fn: Callable) -> float: def _normalize_wx_inject( - wx_inject: Union[bool, Tuple[bool, ...], int], + wx_inject: bool | tuple[bool, ...] | int, num_layers: int, -) -> Tuple[bool, ...]: +) -> tuple[bool, ...]: """Convert wx_inject specification to a boolean tuple. Args: @@ -127,7 +127,7 @@ def __init__( output_dim: int = 1, rectifier_fn: Callable[[jax.Array], jax.Array] = jax.nn.softplus, act_fn: Callable[[jax.Array], jax.Array] = jax.nn.relu, - wx_inject: Union[bool, Tuple[bool, ...], int] = True, + wx_inject: bool | tuple[bool, ...] | int = True, use_bias: bool = True, use_softmax: bool = False, use_sinkhorn: bool = False, @@ -323,15 +323,15 @@ def __init__( dim_hidden: Sequence[int], *, input_dim: int, - output_dim: Optional[int] = None, - num_outputs: Optional[int] = None, + output_dim: int | None = None, + num_outputs: int | None = None, resnet: bool = False, act_fn: Callable[[jax.Array], jax.Array] = jax.nn.relu, - wx_inject: Union[bool, Tuple[bool, ...], int] = True, + wx_inject: bool | tuple[bool, ...] | int = True, use_bias: bool = True, kernel_init: nnx.initializers.Initializer = DEFAULT_KERNEL_INIT, bias_init: nnx.initializers.Initializer = DEFAULT_BIAS_INIT, - final_layer_scale: Optional[float] = None, + final_layer_scale: float | None = None, rngs: nnx.Rngs, ): super().__init__() diff --git a/src/ott/neural/networks/layers/conjugate.py b/src/ott/neural/networks/layers/conjugate.py index 3e1b7e2e0..a90cf78a3 100644 --- a/src/ott/neural/networks/layers/conjugate.py +++ b/src/ott/neural/networks/layers/conjugate.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc -from typing import Callable, Literal, NamedTuple, Optional +from collections.abc import Callable +from typing import Literal, NamedTuple +import jax import jax.numpy as jnp from jaxopt import LBFGS @@ -36,7 +38,7 @@ class ConjugateResults(NamedTuple): num_iter: the number of iterations taken by the solver """ val: float - grad: jnp.ndarray + grad: jax.Array num_iter: int @@ -50,9 +52,9 @@ class FenchelConjugateSolver(abc.ABC): @abc.abstractmethod def solve( self, - f: Callable[[jnp.ndarray], jnp.ndarray], - y: jnp.ndarray, - x_init: Optional[jnp.ndarray] = None + f: Callable[[jax.Array], jax.Array], + y: jax.Array, + x_init: jax.Array | None = None ) -> ConjugateResults: """Solve for the conjugate. @@ -90,9 +92,9 @@ class FenchelConjugateLBFGS(FenchelConjugateSolver): def solve( # noqa: D102 self, - f: Callable[[jnp.ndarray], jnp.ndarray], - y: jnp.ndarray, - x_init: Optional[jnp.array] = None + f: Callable[[jax.Array], jax.Array], + y: jax.Array, + x_init: jnp.array | None = None ) -> ConjugateResults: assert y.ndim diff --git a/src/ott/neural/networks/layers/initializers.py b/src/ott/neural/networks/layers/initializers.py index 3ccca66f4..59e82e0f4 100644 --- a/src/ott/neural/networks/layers/initializers.py +++ b/src/ott/neural/networks/layers/initializers.py @@ -18,16 +18,18 @@ and variance propagation through layers. """ import math -from typing import Callable, Literal, Tuple +from collections.abc import Callable +from typing import Literal import jax import jax.numpy as jnp +from jax.typing import DTypeLike from flax import nnx __all__ = ["get_rectifier_inverse", "principled_icnn_init"] -Initializer = Callable[[jax.Array, Tuple[int, ...], jnp.dtype], jax.Array] +Initializer = Callable[[jax.Array, tuple[int, ...], DTypeLike], jax.Array] RectifierName = Literal["exp", "softplus", "relu", "identity"] @@ -59,7 +61,7 @@ def _principled_icnn_weights( *, alpha: float, rho: float = 0.5, -) -> Tuple[jax.Array, jax.Array, jax.Array]: +) -> tuple[jax.Array, jax.Array, jax.Array]: """Compute log-normal weight parameters for principled ICNN init.""" def _corr_func(fan_in: int, *, rho: float) -> float: @@ -93,7 +95,7 @@ def principled_icnn_init( rectifier_fn: Callable[[jax.Array], jax.Array] = jax.nn.softplus, target_rho: float = 0.5, target_var: float = 1.0, -) -> Tuple[Initializer, Initializer]: +) -> tuple[Initializer, Initializer]: """Compute principled weight and bias initializers for ICNN layers. Implements the initialization from :cite:`hoedt:2023` that @@ -126,8 +128,8 @@ def principled_icnn_init( def weights_init( rng: jax.Array, - shape: Tuple[int, ...], - dtype: jnp.dtype = None + shape: tuple[int, ...], + dtype: DTypeLike | None = None ) -> jax.Array: w = nnx.initializers.normal(stddev=w_log_var ** 0.5)(rng, shape, dtype) w = jnp.exp(w_log_mean + w) @@ -135,8 +137,8 @@ def weights_init( def biases_init( rng: jax.Array, - shape: Tuple[int, ...], - dtype: jnp.dtype = None + shape: tuple[int, ...], + dtype: DTypeLike | None = None ) -> jax.Array: return nnx.initializers.constant(b_mean)(rng, shape, dtype) @@ -148,7 +150,7 @@ def _principled_init_fixed( *, alpha: float, inv_fn: Callable[[jax.Array], jax.Array], -) -> Tuple[Initializer, Initializer]: +) -> tuple[Initializer, Initializer]: """Optimized principled init for target_rho=0.5, target_var=1.0.""" def get_factors(): diff --git a/src/ott/neural/networks/layers/posdef.py b/src/ott/neural/networks/layers/posdef.py index 944482648..7dce6a176 100644 --- a/src/ott/neural/networks/layers/posdef.py +++ b/src/ott/neural/networks/layers/posdef.py @@ -13,7 +13,7 @@ # limitations under the License. """Positive-weight dense layer for input convex neural networks.""" -from typing import Callable, Optional +from collections.abc import Callable import jax import jax.numpy as jnp @@ -80,8 +80,7 @@ def __init__( in_features: int, out_features: int, *, - rectifier_fn: Optional[Callable[[jax.Array], - jax.Array]] = jax.nn.softplus, + rectifier_fn: Callable[[jax.Array], jax.Array] | None = jax.nn.softplus, use_softmax: bool = False, use_sinkhorn: bool = False, use_bias: bool = True, diff --git a/src/ott/neural/networks/potentials.py b/src/ott/neural/networks/potentials.py index ff8a4da4e..2a08a8c75 100644 --- a/src/ott/neural/networks/potentials.py +++ b/src/ott/neural/networks/potentials.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc -from typing import Any, Callable, Optional, Sequence, Tuple, Union +from collections.abc import Callable, Sequence +from typing import Any import jax import jax.numpy as jnp @@ -33,8 +34,8 @@ "LinenMLP", ] -PotentialValueFn_t = Callable[[jnp.ndarray], jnp.ndarray] -PotentialGradientFn_t = Callable[[jnp.ndarray], jnp.ndarray] +PotentialValueFn_t = Callable[[jax.Array], jax.Array] +PotentialGradientFn_t = Callable[[jax.Array], jax.Array] # --------------------------------------------------------------------------- @@ -53,9 +54,9 @@ class PotentialTrainState(train_state.TrainState): potential_gradient_fn: the potential's gradient function """ potential_value_fn: Callable[ - [frozen_dict.FrozenDict[str, jnp.ndarray], Optional[PotentialValueFn_t]], + [frozen_dict.FrozenDict[str, jax.Array], PotentialValueFn_t | None], PotentialValueFn_t] = struct.field(pytree_node=False) - potential_gradient_fn: Callable[[frozen_dict.FrozenDict[str, jnp.ndarray]], + potential_gradient_fn: Callable[[frozen_dict.FrozenDict[str, jax.Array]], PotentialGradientFn_t] = struct.field( pytree_node=False ) @@ -82,8 +83,8 @@ def is_potential(self) -> bool: def potential_value_fn( self, - params: frozen_dict.FrozenDict[str, jnp.ndarray], - other_potential_value_fn: Optional[PotentialValueFn_t] = None, + params: frozen_dict.FrozenDict[str, jax.Array], + other_potential_value_fn: PotentialValueFn_t | None = None, ) -> PotentialValueFn_t: r"""Return a function giving the value of the potential. @@ -114,7 +115,7 @@ def potential_value_fn( "The value of the gradient-based potential depends " \ "on the value of the other potential." - def value_fn(x: jnp.ndarray) -> jnp.ndarray: + def value_fn(x: jax.Array) -> jax.Array: squeeze = x.ndim == 1 if squeeze: x = jnp.expand_dims(x, 0) @@ -127,7 +128,7 @@ def value_fn(x: jnp.ndarray) -> jnp.ndarray: def potential_gradient_fn( self, - params: frozen_dict.FrozenDict[str, jnp.ndarray], + params: frozen_dict.FrozenDict[str, jax.Array], ) -> PotentialGradientFn_t: """Return a function returning a vector or the gradient of the potential. @@ -145,7 +146,7 @@ def create_train_state( self, rng: jax.Array, optimizer: optax.OptState, - input: Union[int, Tuple[int, ...]], + input: int | tuple[int, ...], **kwargs: Any, ) -> PotentialTrainState: """Create initial training state.""" @@ -181,10 +182,10 @@ class LinenPotentialMLP(BasePotential): dim_hidden: Sequence[int] is_potential: bool = True - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.leaky_relu + act_fn: Callable[[jax.Array], jax.Array] = nn.leaky_relu @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # noqa: D102 + def __call__(self, x: jax.Array) -> jax.Array: # noqa: D102 squeeze = x.ndim == 1 if squeeze: x = jnp.expand_dims(x, 0) @@ -221,10 +222,10 @@ class LinenMLP(nn.Module): """ dim_hidden: Sequence[int] - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.elu + act_fn: Callable[[jax.Array], jax.Array] = jax.nn.elu @nn.compact - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + def __call__(self, x: jax.Array) -> jax.Array: """Apply MLP transform.""" for feat in self.dim_hidden[:-1]: x = self.act_fn(nn.Dense(feat)(x)) @@ -250,7 +251,7 @@ def is_potential(self) -> bool: def potential_value_fn( self, - other_potential_value_fn: Optional[PotentialValueFn_t] = None, + other_potential_value_fn: PotentialValueFn_t | None = None, ) -> PotentialValueFn_t: r"""Return a callable giving the potential value. @@ -276,7 +277,7 @@ def potential_value_fn( "on the value of the other potential." ) - def value_fn(x: jnp.ndarray) -> jnp.ndarray: + def value_fn(x: jax.Array) -> jax.Array: squeeze = x.ndim == 1 if squeeze: x = jnp.expand_dims(x, 0) @@ -321,7 +322,7 @@ def __init__( *, input_dim: int, is_potential: bool = True, - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.leaky_relu, + act_fn: Callable[[jax.Array], jax.Array] = jax.nn.leaky_relu, rngs: nnx.Rngs, ): super().__init__() @@ -346,7 +347,7 @@ def __init__( def is_potential(self) -> bool: # noqa: D102 return self._is_potential - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: # noqa: D102 + def __call__(self, x: jax.Array) -> jax.Array: # noqa: D102 squeeze = x.ndim == 1 if squeeze: x = jnp.expand_dims(x, 0) @@ -385,7 +386,7 @@ def __init__( dim_hidden: Sequence[int], *, input_dim: int, - act_fn: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.elu, + act_fn: Callable[[jax.Array], jax.Array] = jax.nn.elu, rngs: nnx.Rngs, ): self._act_fn = act_fn @@ -395,7 +396,7 @@ def __init__( self.layers.append(nnx.Linear(prev_dim, feat, rngs=rngs)) prev_dim = feat - def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + def __call__(self, x: jax.Array) -> jax.Array: """Apply MLP transform.""" for layer in self.layers[:-1]: x = self._act_fn(layer(x)) diff --git a/src/ott/neural/networks/velocity_field/mlp.py b/src/ott/neural/networks/velocity_field/mlp.py index 38eedd63a..0dbf69118 100644 --- a/src/ott/neural/networks/velocity_field/mlp.py +++ b/src/ott/neural/networks/velocity_field/mlp.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Callable, Optional, Sequence +from collections.abc import Callable, Sequence +from typing import Any import jax from jax import numpy as jnp @@ -42,7 +43,7 @@ def __init__( hidden_dims: Sequence[int] = (), cond_dim: int = 0, act_fn: Callable[[jax.Array], jax.Array] = nnx.silu, - time_enc_num_freqs: Optional[int] = None, + time_enc_num_freqs: int | None = None, dropout_rate: float = 0.0, rngs: nnx.Rngs, **kwargs: Any, @@ -77,9 +78,9 @@ def __call__( self, t: jax.Array, x: jax.Array, - cond: Optional[jax.Array] = None, + cond: jax.Array | None = None, *, - rngs: Optional[nnx.Rngs] = None, + rngs: nnx.Rngs | None = None, ) -> jax.Array: """Compute the velocity. @@ -116,7 +117,7 @@ def __init__( self.dropout = nnx.Dropout(dropout_rate) def __call__( - self, x: jax.Array, *, rngs: Optional[nnx.Rngs] = None + self, x: jax.Array, *, rngs: nnx.Rngs | None = None ) -> jax.Array: return self.dropout(self.act_fn(self.lin(x)), rngs=rngs) diff --git a/src/ott/neural/networks/velocity_field/unet.py b/src/ott/neural/networks/velocity_field/unet.py index 1798c5ae7..8fb745d93 100644 --- a/src/ott/neural/networks/velocity_field/unet.py +++ b/src/ott/neural/networks/velocity_field/unet.py @@ -15,10 +15,11 @@ import abc import functools import math -from typing import Any, Literal, Optional, Tuple, Union +from typing import Any, Literal import jax import jax.numpy as jnp +from jax.typing import DTypeLike from flax import nnx @@ -45,21 +46,21 @@ def timestep_embedding( class GroupNorm32(nnx.GroupNorm): def __call__( - self, x: jax.Array, *, mask: Optional[jax.Array] = None + self, x: jax.Array, *, mask: jax.Array | None = None ) -> jax.Array: return super().__call__(x.astype(jnp.float32), mask=mask).astype(x.dtype) def conv_nd( dims: int, - in_channels: Union[int, Tuple[int, ...]], - out_channels: Union[int, Tuple[int, ...]], - kernel_size: Union[int, Tuple[int, ...]], - strides: Union[int, Tuple[int, ...]] = 1, + in_channels: int | tuple[int, ...], + out_channels: int | tuple[int, ...], + kernel_size: int | tuple[int, ...], + strides: int | tuple[int, ...] = 1, *, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, - padding: Union[int, Tuple[int, ...]] = 0, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, + padding: int | tuple[int, ...] = 0, zero_init: bool = False, rngs: nnx.Rngs, **kwargs: Any, @@ -91,8 +92,8 @@ def conv_nd( def normalization( channels: int, *, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, rngs: nnx.Rngs, ) -> nnx.GroupNorm: return GroupNorm32( @@ -114,7 +115,7 @@ def __call__( x: jax.Array, emb: jax.Array, *, - rngs: Optional[nnx.Rngs] = None, + rngs: nnx.Rngs | None = None, ) -> jax.Array: pass @@ -128,9 +129,9 @@ def __init__(self, *layers: nnx.Module): def __call__( self, x: jax.Array, - emb: Optional[jax.Array] = None, + emb: jax.Array | None = None, *, - rngs: Optional[nnx.Rngs] = None, + rngs: nnx.Rngs | None = None, ) -> jax.Array: for layer in self.layers: if isinstance(layer, TimestepBlock): @@ -147,9 +148,9 @@ def __init__( channels: int, use_conv: bool, *, - out_channels: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, + out_channels: int | None = None, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, rngs: nnx.Rngs, ): super().__init__() @@ -186,9 +187,9 @@ def __init__( channels: int, use_conv: bool, *, - out_channels: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, + out_channels: int | None = None, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, rngs: nnx.Rngs, ): super().__init__() @@ -224,12 +225,12 @@ def __init__( emb_channels: int, dropout: float, *, - out_channels: Optional[int] = None, + out_channels: int | None = None, use_conv: bool = False, up: bool = False, down: bool = False, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, rngs: nnx.Rngs, ): super().__init__() @@ -344,7 +345,7 @@ def __call__( x: jax.Array, emb: jax.Array, *, - rngs: Optional[nnx.Rngs] = None, + rngs: nnx.Rngs | None = None, ) -> jax.Array: if self.updown: h = self.in_norm(x) @@ -374,7 +375,7 @@ class QKVAttention(nnx.Module): def __init__( self, n_heads: int, - attn_implementation: Optional[Literal["xla", "cudnn"]] = None, + attn_implementation: Literal["xla", "cudnn"] | None = None, ): super().__init__() self.n_heads = n_heads @@ -410,9 +411,9 @@ def __init__( channels: int, *, num_heads: int = 1, - attn_implementation: Optional[Literal["xla", "cudnn"]] = None, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, + attn_implementation: Literal["xla", "cudnn"] | None = None, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, rngs: nnx.Rngs, ): super().__init__() @@ -488,22 +489,22 @@ class UNet(nnx.Module): def __init__( self, *, - shape: Tuple[int, int, int], + shape: tuple[int, int, int], model_channels: int, num_res_blocks: int, - attention_resolutions: Tuple[int, ...], - out_channels: Optional[int] = None, + attention_resolutions: tuple[int, ...], + out_channels: int | None = None, dropout: float = 0.0, - channel_mult: Tuple[float, ...] = (1, 2, 4, 8), - time_embed_dim: Optional[Union[int, float]] = None, + channel_mult: tuple[float, ...] = (1, 2, 4, 8), + time_embed_dim: int | float | None = None, conv_resample: bool = True, num_heads: int = 1, - num_heads_upsample: Optional[int] = None, + num_heads_upsample: int | None = None, resblock_updown: bool = False, - num_classes: Optional[int] = None, - dtype: Optional[jnp.dtype] = None, - param_dtype: jnp.dtype = jnp.float32, - attn_implementation: Optional[Literal["xla", "cudnn"]] = None, + num_classes: int | None = None, + dtype: DTypeLike | None = None, + param_dtype: DTypeLike = jnp.float32, + attn_implementation: Literal["xla", "cudnn"] | None = None, rngs: nnx.Rngs, ): super().__init__() @@ -742,9 +743,9 @@ def __call__( self, t: jax.Array, x: jax.Array, - cond: Optional[jax.Array] = None, + cond: jax.Array | None = None, *, - rngs: Optional[nnx.Rngs] = None, + rngs: nnx.Rngs | None = None, ) -> jax.Array: """Compute the velocity. diff --git a/src/ott/problems/linear/barycenter_problem.py b/src/ott/problems/linear/barycenter_problem.py index ac466a73c..65a33ab3f 100644 --- a/src/ott/problems/linear/barycenter_problem.py +++ b/src/ott/problems/linear/barycenter_problem.py @@ -11,7 +11,8 @@ # 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 typing import Any, Dict, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import jax import jax.numpy as jnp @@ -50,11 +51,11 @@ class FreeBarycenterProblem: def __init__( self, - y: jnp.ndarray, - b: Optional[jnp.ndarray] = None, - weights: Optional[jnp.ndarray] = None, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, + y: jax.Array, + b: jax.Array | None = None, + weights: jax.Array | None = None, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, **kwargs: Any, ): self._y = y @@ -78,7 +79,7 @@ def __init__( "Point clouds and weights do not have matching shapes." @property - def segmented_y_b(self) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + def segmented_y_b(self) -> tuple[jax.Array, jax.Array, jax.Array]: """Tuple of arrays containing segmented measures, weights, # of points. - Segmented measures of shape ``[num_measures, max_measure_size, ndim]``. @@ -98,20 +99,20 @@ def segmented_y_b(self) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: return y, b, num_per_measure @property - def num_per_measure(self) -> jnp.ndarray: + def num_per_measure(self) -> jax.Array: """``[num_measures,]`` array containing number of points per measure.""" _, _, num_per_measure = self.segmented_y_b return num_per_measure @property - def flattened_y(self) -> jnp.ndarray: + def flattened_y(self) -> jax.Array: """Array of shape ``[num_measures * (N_1 + N_2 + ...), ndim]``.""" if self._is_segmented: return self._y.reshape((-1, self._y.shape[-1])) return self._y @property - def flattened_b(self) -> Optional[jnp.ndarray]: + def flattened_b(self) -> jax.Array | None: """Array of shape ``[num_measures * (N_1 + N_2 + ...),]``.""" return None if self._b is None else self._b.ravel() @@ -131,7 +132,7 @@ def ndim(self) -> int: return self._y.shape[-1] @property - def weights(self) -> jnp.ndarray: + def weights(self) -> jax.Array: """Barycenter weights of shape ``[num_measures,]`` that sum to 1.""" if self._weights is None: return jnp.ones((self.num_measures,)) / self.num_measures @@ -144,7 +145,7 @@ def weights(self) -> jnp.ndarray: def _is_segmented(self) -> bool: return self._y.ndim == 3 - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([self._y, self._b, self._weights], { "cost_fn": self.cost_fn, "epsilon": self.epsilon, @@ -153,7 +154,7 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "FreeBarycenterProblem": y, b, weights = children return cls(y=y, b=b, weights=weights, **aux_data) @@ -175,8 +176,8 @@ class FixedBarycenterProblem: def __init__( self, geom: geometry.Geometry, - a: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, + a: jax.Array, + weights: jax.Array | None = None, ): self.geom = geom self.a = a @@ -188,7 +189,7 @@ def num_measures(self) -> int: return self.a.shape[0] @property - def weights(self) -> jnp.ndarray: + def weights(self) -> jax.Array: """Barycenter weights of shape ``[num_measures,]`` that sum to :math`1`.""" if self._weights is None: return jnp.ones((self.num_measures,)) / self.num_measures @@ -203,7 +204,7 @@ def tree_flatten(self): # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "FixedBarycenterProblem": del aux_data geom, a, weights = children diff --git a/src/ott/problems/linear/linear_problem.py b/src/ott/problems/linear/linear_problem.py index 0e390d0e6..7869157b6 100644 --- a/src/ott/problems/linear/linear_problem.py +++ b/src/ott/problems/linear/linear_problem.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple +from collections.abc import Callable, Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -22,9 +23,8 @@ __all__ = ["LinearProblem"] # TODO(michalk8): move to typing.py when refactoring the types -MarginalFunc = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray] -TransportAppFunc = Callable[[jnp.ndarray, jnp.ndarray, jnp.ndarray, int], - jnp.ndarray] +MarginalFunc = Callable[[jax.Array, jax.Array], jax.Array] +TransportAppFunc = Callable[[jax.Array, jax.Array, jax.Array, int], jax.Array] @jax.tree_util.register_pytree_node_class @@ -51,8 +51,8 @@ class LinearProblem: def __init__( self, geom: geometry.Geometry, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, + b: jax.Array | None = None, tau_a: float = 1.0, tau_b: float = 1.0 ): @@ -63,7 +63,7 @@ def __init__( self.tau_b = tau_b @property - def a(self) -> jnp.ndarray: + def a(self) -> jax.Array: """First marginal.""" if self._a is not None: return self._a @@ -71,7 +71,7 @@ def a(self) -> jnp.ndarray: return jnp.full((n,), fill_value=1.0 / n, dtype=self.dtype) @property - def b(self) -> jnp.ndarray: + def b(self) -> jax.Array: """Second marginal.""" if self._b is not None: return self._b @@ -112,7 +112,7 @@ def potential_fn_from_dual_vec( self, fg: jax.Array, *, - epsilon: Optional[float] = None, + epsilon: float | None = None, axis: Literal[0, 1], ) -> Callable[[jax.Array], jax.Array]: r"""Get potential function from a dual vector using the :term:`c-transform`. @@ -150,16 +150,16 @@ def _c_transform( self, fg: jax.Array, *, - epsilon: Optional[float] = None, + epsilon: float | None = None, axis: Literal[0, 1], - ) -> Tuple[jax.Array, jax.Array]: + ) -> tuple[jax.Array, jax.Array]: - def _soft_c_transform(fg: jax.Array) -> Tuple[jax.Array, jax.Array]: + def _soft_c_transform(fg: jax.Array) -> tuple[jax.Array, jax.Array]: cost = self.geom.cost_matrix z = (fg - cost) / epsilon return -epsilon * math_utils.logsumexp(z, b=self.b, axis=axis), z - def _hard_c_transform(fg: jax.Array) -> Tuple[jax.Array, jax.Array]: + def _hard_c_transform(fg: jax.Array) -> tuple[jax.Array, jax.Array]: cost = self.geom.cost_matrix z = fg - cost pos_weights = self.b[None, :] > 0.0 @@ -172,7 +172,7 @@ def _hard_c_transform(fg: jax.Array) -> Tuple[jax.Array, jax.Array]: def get_transport_functions( self, lse_mode: bool - ) -> Tuple[MarginalFunc, MarginalFunc, TransportAppFunc]: + ) -> tuple[MarginalFunc, MarginalFunc, TransportAppFunc]: """Instantiate useful functions for Sinkhorn depending on lse_mode.""" geom = self.geom if lse_mode: @@ -192,7 +192,7 @@ def get_transport_functions( ) return marginal_a, marginal_b, app_transport - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([self.geom, self._a, self._b], { "tau_a": self.tau_a, "tau_b": self.tau_b @@ -200,6 +200,6 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "LinearProblem": return cls(*children, **aux_data) diff --git a/src/ott/problems/linear/potentials.py b/src/ott/problems/linear/potentials.py index 8b23627f6..060da8662 100644 --- a/src/ott/problems/linear/potentials.py +++ b/src/ott/problems/linear/potentials.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Any, Callable, Dict, Optional, Tuple +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -45,11 +46,11 @@ class DualPotentials: g: The second dual potential function. cost_fn: The cost function used to solve the OT problem. """ - f: Optional[PotentialFn] - g: Optional[PotentialFn] + f: PotentialFn | None + g: PotentialFn | None cost_fn: costs.CostFn - def transport(self, vec: jnp.ndarray, forward: bool = True) -> jnp.ndarray: + def transport(self, vec: jax.Array, forward: bool = True) -> jax.Array: r"""Transport ``vec`` according to Gangbo-McCann Brenier :cite:`brenier:91`. Uses Proposition 1.15 from :cite:`santambrogio:15` to compute an OT map when @@ -83,7 +84,7 @@ def transport(self, vec: jnp.ndarray, forward: bool = True) -> jnp.ndarray: return twist_op(vec, self._grad_f(vec), False) return twist_op(vec, self._grad_g(vec), True) - def distance(self, src: jnp.ndarray, tgt: jnp.ndarray) -> float: + def distance(self, src: jax.Array, tgt: jax.Array) -> float: r"""Evaluate Wasserstein distance between samples using dual potentials. This uses direct estimation of potentials against measures when dual @@ -102,27 +103,27 @@ def distance(self, src: jnp.ndarray, tgt: jnp.ndarray) -> float: return jnp.mean(f(src)) + jnp.mean(g(tgt)) @property - def _grad_f(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + def _grad_f(self) -> Callable[[jax.Array], jax.Array]: """Vectorized gradient of the potential function :attr:`f`.""" assert self.f is not None, "The `f` potential is not computed." return jax.vmap(jax.grad(self.f, argnums=0)) @property - def _grad_g(self) -> Callable[[jnp.ndarray], jnp.ndarray]: + def _grad_g(self) -> Callable[[jax.Array], jax.Array]: """Vectorized gradient of the potential function :attr:`g`.""" assert self.g is not None, "The `g` potential is not computed." return jax.vmap(jax.grad(self.g, argnums=0)) def plot_ot_map( self, - source: jnp.ndarray, - target: jnp.ndarray, - samples: Optional[jnp.ndarray] = None, + source: jax.Array, + target: jax.Array, + samples: jax.Array | None = None, forward: bool = True, - ax: Optional["plt.Axes"] = None, - scatter_kwargs: Optional[Dict[str, Any]] = None, - legend_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple["plt.Figure", "plt.Axes"]: + ax: "plt.Axes | None" = None, + scatter_kwargs: dict[str, Any] | None = None, + legend_kwargs: dict[str, Any] | None = None, + ) -> tuple["plt.Figure", "plt.Axes"]: """Plot data and learned optimal transport map. Args: @@ -211,12 +212,12 @@ def plot_potential( forward: bool = True, quantile: float = 0.05, kantorovich: bool = True, - ax: Optional["mpl.axes.Axes"] = None, - x_bounds: Tuple[float, float] = (-6, 6), - y_bounds: Tuple[float, float] = (-6, 6), + ax: "mpl.axes.Axes | None" = None, + x_bounds: tuple[float, float] = (-6, 6), + y_bounds: tuple[float, float] = (-6, 6), num_grid: int = 50, - contourf_kwargs: Optional[Dict[str, Any]] = None, - ) -> Tuple["mpl.figure.Figure", "mpl.axes.Axes"]: + contourf_kwargs: dict[str, Any] | None = None, + ) -> tuple["mpl.figure.Figure", "mpl.axes.Axes"]: r"""Plot the potential. Args: diff --git a/src/ott/problems/linear/semidiscrete_linear_problem.py b/src/ott/problems/linear/semidiscrete_linear_problem.py index 9b8b61d1f..ed862e8ef 100644 --- a/src/ott/problems/linear/semidiscrete_linear_problem.py +++ b/src/ott/problems/linear/semidiscrete_linear_problem.py @@ -11,7 +11,7 @@ # 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 typing import Callable, Optional +from collections.abc import Callable import jax import jax.numpy as jnp @@ -39,7 +39,7 @@ class SemidiscreteLinearProblem: def __init__( self, geom: semidiscrete_pointcloud.SemidiscretePointCloud, - b: Optional[jax.Array] = None, + b: jax.Array | None = None, tau_b: float = 1.0, ): assert tau_b == 1.0, "Unbalanced semidiscrete problem is not supported." @@ -52,7 +52,7 @@ def sample( rng: jax.Array, num_samples: int, *, - epsilon: Optional[float] = None, + epsilon: float | None = None, ) -> linear_problem.LinearProblem: """Sample a linear OT problem. @@ -75,7 +75,7 @@ def potential_fn_from_dual_vec( self, g: jax.Array, *, - epsilon: Optional[float] = None + epsilon: float | None = None ) -> Callable[[jax.Array], jax.Array]: r"""Get potential function from a dual vector using the :term:`c-transform`. @@ -92,7 +92,7 @@ def potential_fn_from_dual_vec( return prob.potential_fn_from_dual_vec(g, epsilon=epsilon, axis=1) @property - def b(self) -> jnp.ndarray: + def b(self) -> jax.Array: """Second marginal.""" if self._b is not None: return self._b diff --git a/src/ott/problems/quadratic/gw_barycenter.py b/src/ott/problems/quadratic/gw_barycenter.py index 263f48506..a8c0b573b 100644 --- a/src/ott/problems/quadratic/gw_barycenter.py +++ b/src/ott/problems/quadratic/gw_barycenter.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Dict, Literal, Optional, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -60,14 +61,14 @@ class GWBarycenterProblem(barycenter_problem.FreeBarycenterProblem): def __init__( self, - y: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - weights: Optional[jnp.ndarray] = None, - costs: Optional[jnp.ndarray] = None, - y_fused: Optional[jnp.ndarray] = None, + y: jax.Array | None = None, + b: jax.Array | None = None, + weights: jax.Array | None = None, + costs: jax.Array | None = None, + y_fused: jax.Array | None = None, fused_penalty: float = 1.0, gw_loss: Literal["sqeucl", "kl"] = "sqeucl", - scale_cost: Union[float, Literal["mean", "max_cost"]] = 1.0, + scale_cost: float | Literal["mean", "max_cost"] = 1.0, **kwargs: Any, ): assert y is None or costs is None, "Cannot specify both `y` and `costs`." @@ -98,9 +99,7 @@ def __init__( # TODO(michalk8): in the future, consider checking the other 2 cases # using `segmented_y` and `segmented_y_fused`? - def update_barycenter( - self, transports: jnp.ndarray, a: jnp.ndarray - ) -> jnp.ndarray: + def update_barycenter(self, transports: jax.Array, a: jax.Array) -> jax.Array: """Update the barycenter cost matrix. Uses the eq. 14 and 15 of :cite:`peyre:16`. @@ -116,11 +115,11 @@ def update_barycenter( @functools.partial(jax.vmap, in_axes=[0, 0, 0, None]) def project( - y: jnp.ndarray, - b: jnp.ndarray, - transport: jnp.ndarray, - fn: Optional[quadratic_costs.Loss], - ) -> jnp.ndarray: + y: jax.Array, + b: jax.Array, + transport: jax.Array, + fn: quadratic_costs.Loss | None, + ) -> jax.Array: geom = self._create_y_geometry(y) fn, lin = (None, True) if fn is None else (fn.func, fn.is_linear) @@ -146,8 +145,9 @@ def project( return jnp.exp(barycenter) return barycenter - def update_features(self, transports: jnp.ndarray, - a: jnp.ndarray) -> Optional[jnp.ndarray]: + def update_features( + self, transports: jax.Array, a: jax.Array + ) -> jax.Array | None: """Update the barycenter features in the fused case :cite:`vayer:19`. Uses :cite:`cuturi:14` eq. 8, and is implemented only @@ -181,7 +181,7 @@ def update_features(self, transports: jnp.ndarray, def _create_bary_geometry( self, - cost_matrix: jnp.ndarray, + cost_matrix: jax.Array, ) -> geometry.Geometry: return geometry.Geometry( cost_matrix=cost_matrix, @@ -191,7 +191,7 @@ def _create_bary_geometry( def _create_y_geometry( self, - y: jnp.ndarray, + y: jax.Array, ) -> geometry.Geometry: if self._y_as_costs: assert y.shape[0] == y.shape[1], y.shape @@ -209,8 +209,8 @@ def _create_y_geometry( def _create_fused_geometry( self, - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, ) -> pointcloud.PointCloud: return pointcloud.PointCloud( x, @@ -223,9 +223,9 @@ def _create_fused_geometry( def _create_problem( self, state: "GWBarycenterState", # noqa: F821 - y: jnp.ndarray, - b: jnp.ndarray, - f: Optional[jnp.ndarray] = None + y: jax.Array, + b: jax.Array, + f: jax.Array | None = None ) -> quadratic_problem.QuadraticProblem: geom_xx = self._create_bary_geometry(state.cost) geom_yy = self._create_y_geometry(y) @@ -251,7 +251,7 @@ def is_fused(self) -> bool: return self._y_fused is not None @property - def segmented_y_fused(self) -> Optional[jnp.ndarray]: + def segmented_y_fused(self) -> jax.Array | None: """Feature array of shape used in the fused case.""" if not self.is_fused or self._y_fused.ndim == 3: return self._y_fused @@ -263,11 +263,11 @@ def segmented_y_fused(self) -> Optional[jnp.ndarray]: return y_fused @property - def ndim(self) -> Optional[int]: # noqa: D102 + def ndim(self) -> int | None: # noqa: D102 return None if self._y_as_costs else self._y.shape[-1] @property - def ndim_fused(self) -> Optional[int]: + def ndim_fused(self) -> int | None: """Number of dimensions of the fused term.""" return self._y_fused.shape[-1] if self.is_fused else None @@ -286,7 +286,7 @@ def gw_loss(self) -> quadratic_costs.GWLoss: f"Loss `{self._loss_name}` is not yet implemented." ) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 (y, b, weights), aux = super().tree_flatten() if self._y_as_costs: children = [None, b, weights, y] @@ -299,7 +299,7 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "GWBarycenterProblem": y, b, weights, costs, y_fused = children return cls( diff --git a/src/ott/problems/quadratic/quadratic_costs.py b/src/ott/problems/quadratic/quadratic_costs.py index 33a6d504e..83f436081 100644 --- a/src/ott/problems/quadratic/quadratic_costs.py +++ b/src/ott/problems/quadratic/quadratic_costs.py @@ -11,8 +11,10 @@ # 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 typing import Callable, NamedTuple +from collections.abc import Callable +from typing import NamedTuple +import jax import jax.numpy as jnp import jax.scipy as jsp @@ -20,7 +22,7 @@ class Loss(NamedTuple): # noqa: D101 - func: Callable[[jnp.ndarray], jnp.ndarray] + func: Callable[[jax.Array], jax.Array] is_linear: bool diff --git a/src/ott/problems/quadratic/quadratic_problem.py b/src/ott/problems/quadratic/quadratic_problem.py index fe02d8a61..d00bee7ff 100644 --- a/src/ott/problems/quadratic/quadratic_problem.py +++ b/src/ott/problems/quadratic/quadratic_problem.py @@ -11,7 +11,7 @@ # 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 typing import TYPE_CHECKING, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Literal import jax import jax.numpy as jnp @@ -83,17 +83,17 @@ def __init__( self, geom_xx: geometry.Geometry, geom_yy: geometry.Geometry, - geom_xy: Optional[geometry.Geometry] = None, + geom_xy: geometry.Geometry | None = None, fused_penalty: float = 1.0, - scale_cost: Optional[Union[float, str]] = None, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", + scale_cost: float | str | None = None, + a: jax.Array | None = None, + b: jax.Array | None = None, + loss: Literal["sqeucl", "kl"] | quadratic_costs.GWLoss = "sqeucl", tau_a: float = 1.0, tau_b: float = 1.0, gw_unbalanced_correction: bool = True, - ranks: Union[int, Tuple[int, ...]] = -1, - tolerances: Union[float, Tuple[float, ...]] = 1e-2, + ranks: int | tuple[int, ...] = -1, + tolerances: float | tuple[float, ...] = 1e-2, ): if scale_cost is not None: geom_xx = geom_xx.set_scale_cost(scale_cost) @@ -124,8 +124,8 @@ def __init__( def marginal_dependent_cost( self, - marginal_1: jnp.ndarray, - marginal_2: jnp.ndarray, + marginal_1: jax.Array, + marginal_2: jax.Array, ) -> low_rank.LRCGeometry: r"""Initialize cost term that depends on the marginals of the transport. @@ -168,9 +168,9 @@ def marginal_dependent_cost( def cost_unbalanced_correction( self, - transport_matrix: jnp.ndarray, - marginal_1: jnp.ndarray, - marginal_2: jnp.ndarray, + transport_matrix: jax.Array, + marginal_1: jax.Array, + marginal_2: jax.Array, epsilon: float, ) -> float: r"""Calculate cost term from the quadratic divergence when unbalanced. @@ -192,10 +192,10 @@ def cost_unbalanced_correction( :math:`+ epsilon * \sum(KL(P|ab'))` Args: - transport_matrix: jnp.ndarray[num_a, num_b], transport matrix. - marginal_1: jnp.ndarray[num_a,], marginal of the transport matrix + transport_matrix: jax.Array[num_a, num_b], transport matrix. + marginal_1: jax.Array[num_a,], marginal of the transport matrix for samples from :attr:`geom_xx`. - marginal_2: jnp.ndarray[num_b,], marginal of the transport matrix + marginal_2: jax.Array[num_b,], marginal of the transport matrix for samples from :attr:`geom_yy`. epsilon: entropy regularizer. @@ -234,7 +234,7 @@ def init_transport_mass(self) -> float: def update_lr_geom( self, lr_sink: "sinkhorn_lr.LRSinkhornOutput", - relative_epsilon: Optional[Literal["mean", "std"]] = None, + relative_epsilon: Literal["mean", "std"] | None = None, ) -> geometry.Geometry: """Recompute (possibly LRC) linearization using LR Sinkhorn output.""" marginal_1 = lr_sink.marginal(1) @@ -268,9 +268,9 @@ def update_lr_geom( def update_linearization( self, transport: Transport, - epsilon: Optional[float] = None, + epsilon: float | None = None, old_transport_mass: float = 1.0, - relative_epsilon: Optional[Literal["mean", "std"]] = None, + relative_epsilon: Literal["mean", "std"] | None = None, ) -> linear_problem.LinearProblem: """Update linearization of GW problem by updating cost matrix. @@ -337,7 +337,7 @@ def update_lr_linearization( self, lr_sink: "sinkhorn_lr.LRSinkhornOutput", *, - relative_epsilon: Optional[Literal["mean", "std"]] = None, + relative_epsilon: Literal["mean", "std"] | None = None, ) -> linear_problem.LinearProblem: """Update a Quad problem linearization using a LR Sinkhorn.""" return linear_problem.LinearProblem( @@ -349,7 +349,7 @@ def update_lr_linearization( ) @property - def _fused_cost_matrix(self) -> Union[float, jnp.ndarray]: + def _fused_cost_matrix(self) -> float | jax.Array: return self.geom_xy.cost_matrix if self.is_fused else 0.0 @property @@ -372,7 +372,7 @@ def convertible(geom: geometry.Geometry) -> bool: def to_low_rank( self, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, ) -> "QuadraticProblem": """Convert geometries to low-rank. @@ -384,8 +384,8 @@ def to_low_rank( """ def convert( - vals: Union[int, float, Tuple[Union[int, float], ...]] - ) -> Tuple[Union[int, float], ...]: + vals: int | float | tuple[int | float, ...] + ) -> tuple[int | float, ...]: size = 2 + self.is_fused if isinstance(vals, (int, float)): return (vals,) * 3 @@ -425,18 +425,18 @@ def geom_yy(self) -> geometry.Geometry: return self._geom_yy @property - def geom_xy(self) -> Optional[geometry.Geometry]: + def geom_xy(self) -> geometry.Geometry | None: """Geometry of the joint space.""" return self._geom_xy @property - def a(self) -> jnp.ndarray: + def a(self) -> jax.Array: """First marginal.""" num_a = self.geom_xx.shape[0] return jnp.ones((num_a,)) / num_a if self._a is None else self._a @property - def b(self) -> jnp.ndarray: + def b(self) -> jax.Array: """Second marginal.""" num_b = self.geom_yy.shape[0] return jnp.ones((num_b,)) / num_b if self._b is None else self._b @@ -456,12 +456,12 @@ def is_low_rank(self) -> bool: ) @property - def linear_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: + def linear_loss(self) -> tuple[quadratic_costs.Loss, quadratic_costs.Loss]: """Linear part of the Gromov-Wasserstein loss.""" return self.loss.f1, self.loss.f2 @property - def quad_loss(self) -> Tuple[quadratic_costs.Loss, quadratic_costs.Loss]: + def quad_loss(self) -> tuple[quadratic_costs.Loss, quadratic_costs.Loss]: """Quadratic part of the Gromov-Wasserstein loss.""" return self.loss.h1, self.loss.h2 @@ -490,7 +490,7 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 def apply_cost( # noqa: D103 - geom: geometry.Geometry, arr: jnp.ndarray, *, axis: int, + geom: geometry.Geometry, arr: jax.Array, *, axis: int, fn: quadratic_costs.Loss -) -> jnp.ndarray: +) -> jax.Array: return geom.apply_cost(arr, axis=axis, fn=fn.func, is_linear=fn.is_linear) diff --git a/src/ott/solvers/linear/_solve.py b/src/ott/solvers/linear/_solve.py index e5c34d01b..17edb148f 100644 --- a/src/ott/solvers/linear/_solve.py +++ b/src/ott/solvers/linear/_solve.py @@ -11,10 +11,9 @@ # 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 typing import Any, Optional, Union +from typing import Any import jax -import jax.numpy as jnp from ott import utils from ott.geometry import geometry, pointcloud @@ -28,13 +27,13 @@ def solve( geom: geometry.Geometry, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, + b: jax.Array | None = None, tau_a: float = 1.0, tau_b: float = 1.0, rank: int = -1, **kwargs: Any -) -> Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput]: +) -> sinkhorn.SinkhornOutput | sinkhorn_lr.LRSinkhornOutput: """Solve linear regularized OT problem using Sinkhorn iterations. Args: @@ -66,8 +65,8 @@ def solve( def solve_univariate( geom: pointcloud.PointCloud, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, + b: jax.Array | None = None, *, return_transport: bool = False, return_dual_variables: bool = False, @@ -104,8 +103,8 @@ def solve_univariate( def solve_semidiscrete( geom: sdpc.SemidiscretePointCloud, - b: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, + b: jax.Array | None = None, + rng: jax.Array | None = None, **kwargs: Any, ) -> semidiscrete.SemidiscreteOutput: """Solve a (regularized) semidiscrete OT problem. diff --git a/src/ott/solvers/linear/acceleration.py b/src/ott/solvers/linear/acceleration.py index c992e1e52..048b9ac6f 100644 --- a/src/ott/solvers/linear/acceleration.py +++ b/src/ott/solvers/linear/acceleration.py @@ -34,7 +34,7 @@ class AndersonAcceleration: refresh_every: int = 1 # Recompute interpolation periodically. ridge_identity: float = 1e-2 # Ridge used in the linear system. - def extrapolation(self, xs: jnp.ndarray, fxs: jnp.ndarray) -> jnp.ndarray: + def extrapolation(self, xs: jax.Array, fxs: jax.Array) -> jax.Array: """Compute Anderson extrapolation from past observations.""" # Remove -inf values to instantiate quadratic problem. All others # remain since they might be caused by a valid issue. @@ -161,10 +161,10 @@ def lehmann(self, state: "sinkhorn.SinkhornState") -> float: def __call__( # noqa: D102 self, weight: float, - value: jnp.ndarray, - new_value: jnp.ndarray, + value: jax.Array, + new_value: jax.Array, lse_mode: bool = True - ) -> jnp.ndarray: + ) -> jax.Array: if lse_mode: value = jnp.where(jnp.isfinite(value), value, 0.0) return (1.0 - weight) * value + weight * new_value diff --git a/src/ott/solvers/linear/continuous_barycenter.py b/src/ott/solvers/linear/continuous_barycenter.py index 84dc46cc1..c1e62d865 100644 --- a/src/ott/solvers/linear/continuous_barycenter.py +++ b/src/ott/solvers/linear/continuous_barycenter.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, NamedTuple, Optional, Tuple, Union +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -27,7 +27,7 @@ __all__ = ["FreeBarycenterState", "FreeWassersteinBarycenter"] -LinearOutput = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput] +LinearOutput = sinkhorn.SinkhornOutput | sinkhorn_lr.LRSinkhornOutput class FreeBarycenterState(NamedTuple): @@ -46,12 +46,12 @@ class FreeBarycenterState(NamedTuple): at each iteration. """ - x: jnp.ndarray - a: jnp.ndarray - costs: Optional[jnp.ndarray] = None - linear_convergence: Optional[jnp.ndarray] = None - linear_outputs: Optional[LinearOutput] = None - errors: Optional[jnp.ndarray] = None + x: jax.Array + a: jax.Array + costs: jax.Array | None = None + linear_convergence: jax.Array | None = None + linear_outputs: LinearOutput | None = None + errors: jax.Array | None = None def set(self, **kwargs: Any) -> "FreeBarycenterState": """Return a copy of self, possibly with overwrites.""" @@ -76,7 +76,7 @@ def update( @functools.partial(jax.vmap, in_axes=[None, None, 0, 0]) def solve_linear_ot( - a: Optional[jnp.ndarray], x: jnp.ndarray, b: jnp.ndarray, y: jnp.ndarray + a: jax.Array | None, x: jax.Array, b: jax.Array, y: jax.Array ): geom = pointcloud.PointCloud( x, y, cost_fn=bar_prob.cost_fn, epsilon=bar_prob.epsilon @@ -140,20 +140,20 @@ class FreeBarycenterOutput(NamedTuple): at each iteration. """ - x: jnp.ndarray - a: jnp.ndarray + x: jax.Array + a: jax.Array bar_prob: barycenter_problem.FreeBarycenterProblem - costs: jnp.ndarray - linear_convergence: jnp.ndarray + costs: jax.Array + linear_convergence: jax.Array linear_outputs: LinearOutput - errors: Optional[jnp.ndarray] = None + errors: jax.Array | None = None @property def all_linear_solvers_converged(self) -> bool: """Whether all linear convergence flags converged.""" return jnp.all(self.linear_convergence[self.linear_convergence != -1]) - def matrix_at_index(self, measure_index: int) -> jnp.ndarray: + def matrix_at_index(self, measure_index: int) -> jax.Array: """Return the transport matrix from barycenter to measure_index measure.""" size_measure = self.bar_prob.num_per_measure[measure_index] matrix = self.linear_output_at_index(measure_index).matrix @@ -170,7 +170,7 @@ def num_iters(self) -> int: return jnp.sum(self.linear_convergence != -1) @property - def costs_along_iterations(self) -> jnp.ndarray: + def costs_along_iterations(self) -> jax.Array: """Costs vector with superfluous values removed.""" return self.costs[:self.num_iters] @@ -187,8 +187,8 @@ def __call__( # noqa: D102 self, bar_prob: barycenter_problem.FreeBarycenterProblem, bar_size: int = 100, - x_init: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, + x_init: jax.Array | None = None, + rng: jax.Array | None = None, ) -> FreeBarycenterState: rng = utils.default_prng_key(rng) return self.iterations(bar_size, bar_prob, x_init, rng) @@ -197,8 +197,8 @@ def init_state( self, bar_prob: barycenter_problem.FreeBarycenterProblem, bar_size: int, - x_init: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, + x_init: jax.Array | None = None, + rng: jax.Array | None = None, ) -> FreeBarycenterState: """Initialize the state of the Wasserstein barycenter iterations. @@ -279,20 +279,20 @@ def output_from_state( # noqa: D102 def iterations( self, bar_size: int, bar_prob: barycenter_problem.FreeBarycenterProblem, - x_init: jnp.ndarray, rng: jax.Array + x_init: jax.Array, rng: jax.Array ) -> FreeBarycenterState: """Wasserstein barycenter outer loop.""" def cond_fn( iteration: int, - constants: Tuple[FreeWassersteinBarycenter, + constants: tuple[FreeWassersteinBarycenter, barycenter_problem.FreeBarycenterProblem], state: FreeBarycenterState ) -> bool: return self._continue(state, iteration) def body_fn( - iteration, constants: Tuple[FreeWassersteinBarycenter, + iteration, constants: tuple[FreeWassersteinBarycenter, barycenter_problem.FreeBarycenterProblem], state: FreeBarycenterState, compute_error: bool ) -> FreeBarycenterState: diff --git a/src/ott/solvers/linear/discrete_barycenter.py b/src/ott/solvers/linear/discrete_barycenter.py index 10b4e69cc..0d56007ad 100644 --- a/src/ott/solvers/linear/discrete_barycenter.py +++ b/src/ott/solvers/linear/discrete_barycenter.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import NamedTuple, Optional, Sequence +from collections.abc import Sequence +from typing import NamedTuple import jax import jax.numpy as jnp @@ -26,10 +27,10 @@ class SinkhornBarycenterOutput(NamedTuple): # noqa: D101 - f: jnp.ndarray - g: jnp.ndarray - histogram: jnp.ndarray - errors: jnp.ndarray + f: jax.Array + g: jax.Array + histogram: jax.Array + errors: jax.Array @jax.tree_util.register_pytree_node_class @@ -79,7 +80,7 @@ def __init__( def __call__( self, fixed_bp: barycenter_problem.FixedBarycenterProblem, - dual_initialization: Optional[jnp.ndarray] = None, + dual_initialization: jax.Array | None = None, ) -> SinkhornBarycenterOutput: """Solve barycenter problem, possibly using clever initialization. @@ -129,10 +130,10 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 @functools.partial(jax.jit, static_argnums=(5, 6, 7, 8, 9, 10, 11, 12)) def _discrete_barycenter( - geom: geometry.Geometry, a: jnp.ndarray, weights: jnp.ndarray, - dual_initialization: jnp.ndarray, threshold: float, - norm_error: Sequence[int], inner_iterations: int, min_iterations: int, - max_iterations: int, lse_mode: bool, debiased: bool, num_a: int, num_b: int + geom: geometry.Geometry, a: jax.Array, weights: jax.Array, + dual_initialization: jax.Array, threshold: float, norm_error: Sequence[int], + inner_iterations: int, min_iterations: int, max_iterations: int, + lse_mode: bool, debiased: bool, num_a: int, num_b: int ) -> SinkhornBarycenterOutput: """Jit'able function to compute discrete barycenters.""" if lse_mode: diff --git a/src/ott/solvers/linear/implicit_differentiation.py b/src/ott/solvers/linear/implicit_differentiation.py index ef3633714..595ad14d9 100644 --- a/src/ott/solvers/linear/implicit_differentiation.py +++ b/src/ott/solvers/linear/implicit_differentiation.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple +from collections.abc import Callable +from typing import TYPE_CHECKING, Any import jax import jax.numpy as jnp @@ -24,9 +25,8 @@ if TYPE_CHECKING: from ott.problems.linear import linear_problem -LinOp_t = Callable[[jnp.ndarray], jnp.ndarray] -Solver_t = Callable[[LinOp_t, jnp.ndarray, Optional[LinOp_t], bool], - jnp.ndarray] +LinOp_t = Callable[[jax.Array], jax.Array] +Solver_t = Callable[[LinOp_t, jax.Array, LinOp_t | None, bool], jax.Array] __all__ = ["ImplicitDiff", "solve_jax_cg"] @@ -72,19 +72,19 @@ class ImplicitDiff: :term`implicit differentiation`. """ - solver: Optional[Solver_t] = None - solver_kwargs: Optional[Dict[str, Any]] = None + solver: Solver_t | None = None + solver_kwargs: dict[str, Any] | None = None symmetric: bool = False - precondition_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None + precondition_fun: Callable[[jax.Array], jax.Array] | None = None def solve( self, - gr: Tuple[jnp.ndarray, jnp.ndarray], + gr: tuple[jax.Array, jax.Array], ot_prob: "linear_problem.LinearProblem", - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, lse_mode: bool, - ) -> jnp.ndarray: + ) -> jax.Array: r"""Apply minus inverse of Hessian of ``reg_ot_cost`` w.r.t. [``f``, ``g``]. This function is used to carry out :term:`implicit differentiation` of @@ -231,7 +231,7 @@ def solve( return jnp.concatenate((-vjp_gr_f, -vjp_gr_g)) def first_order_conditions( - self, prob, f: jnp.ndarray, g: jnp.ndarray, lse_mode: bool + self, prob, f: jax.Array, g: jax.Array, lse_mode: bool ): r"""Compute vector of first order conditions for the reg-OT problem. @@ -245,12 +245,12 @@ def first_order_conditions( Args: prob: definition of the linear optimal transport problem. - f: jnp.ndarray, first potential - g: jnp.ndarray, second potential + f: jax.Array, first potential + g: jax.Array, second potential lse_mode: bool Returns: - a jnp.ndarray of size (size of ``n + m``) quantifying deviation to + a jax.Array of size (size of ``n + m``) quantifying deviation to optimality for variables ``f`` and ``g``. """ geom = prob.geom @@ -273,8 +273,8 @@ def first_order_conditions( return jnp.concatenate((result_a, result_b)) def gradient( - self, prob: "linear_problem.LinearProblem", f: jnp.ndarray, - g: jnp.ndarray, lse_mode: bool, gr: Tuple[jnp.ndarray, jnp.ndarray] + self, prob: "linear_problem.LinearProblem", f: jax.Array, g: jax.Array, + lse_mode: bool, gr: tuple[jax.Array, jax.Array] ) -> "linear_problem.LinearProblem": """Apply VJP to recover gradient in reverse mode differentiation.""" # Applies first part of vjp to gr: inverse part of implicit function theorem @@ -294,13 +294,13 @@ def replace(self, **kwargs: Any) -> "ImplicitDiff": # noqa: D102 def solve_jax_cg( lin: LinOp_t, - b: jnp.ndarray, - lin_t: Optional[LinOp_t] = None, + b: jax.Array, + lin_t: LinOp_t | None = None, symmetric: bool = False, ridge_identity: float = 0.0, ridge_kernel: float = 0.0, **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Wrapper around JAX native linear solvers. Args: diff --git a/src/ott/solvers/linear/lineax_implicit.py b/src/ott/solvers/linear/lineax_implicit.py index 22e5b116a..90e2382d3 100644 --- a/src/ott/solvers/linear/lineax_implicit.py +++ b/src/ott/solvers/linear/lineax_implicit.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Optional +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp @@ -20,13 +21,13 @@ def _cg( - matvec: Callable[[jnp.ndarray], jnp.ndarray], - b: jnp.ndarray, + matvec: Callable[[jax.Array], jax.Array], + b: jax.Array, *, rtol: float = 1e-6, atol: float = 1e-6, - maxiter: Optional[int] = None, -) -> jnp.ndarray: + maxiter: int | None = None, +) -> jax.Array: """Conjugate gradient solver using jax.lax.while_loop.""" if maxiter is None: maxiter = 10 * b.shape[0] @@ -60,14 +61,14 @@ def body_fun(state): def solve_lineax( lin: Callable, - b: jnp.ndarray, - lin_t: Optional[Callable] = None, + b: jax.Array, + lin_t: Callable | None = None, symmetric: bool = False, - nonsym_solver: Optional[Any] = None, + nonsym_solver: Any | None = None, ridge_identity: float = 0.0, ridge_kernel: float = 0.0, **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Solve a linear system using conjugate gradients. This implementation uses a JAX-native CG solver that works correctly inside diff --git a/src/ott/solvers/linear/lr_utils.py b/src/ott/solvers/linear/lr_utils.py index e7695b02a..1a75a452e 100644 --- a/src/ott/solvers/linear/lr_utils.py +++ b/src/ott/solvers/linear/lr_utils.py @@ -11,7 +11,7 @@ # 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 typing import NamedTuple, Optional, Tuple +from typing import NamedTuple import jax import jax.numpy as jnp @@ -24,27 +24,27 @@ class State(NamedTuple): # noqa: D101 - v1: jnp.ndarray - v2: jnp.ndarray - u1: jnp.ndarray - u2: jnp.ndarray - g: jnp.ndarray + v1: jax.Array + v2: jax.Array + u1: jax.Array + u2: jax.Array + g: jax.Array err: float class Constants(NamedTuple): # noqa: D101 - a: jnp.ndarray - b: jnp.ndarray + a: jax.Array + b: jax.Array rho_a: float rho_b: float - supp_a: Optional[jnp.ndarray] = None - supp_b: Optional[jnp.ndarray] = None + supp_a: jax.Array | None = None + supp_b: jax.Array | None = None def unbalanced_dykstra_lse( - c_q: jnp.ndarray, - c_r: jnp.ndarray, - c_g: jnp.ndarray, + c_q: jax.Array, + c_r: jax.Array, + c_g: jax.Array, gamma: float, ot_prob: linear_problem.LinearProblem, translation_invariant: bool = True, @@ -52,7 +52,7 @@ def unbalanced_dykstra_lse( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array, jax.Array]: """Dykstra's algorithm for the unbalanced :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in LSE mode. @@ -74,10 +74,10 @@ def unbalanced_dykstra_lse( """ # noqa: D205 def _softm( - v: jnp.ndarray, - c: jnp.ndarray, + v: jax.Array, + c: jax.Array, axis: int, - ) -> jnp.ndarray: + ) -> jax.Array: v = jnp.expand_dims(v, axis=1 - axis) return jsp.special.logsumexp(v + c, axis=axis) @@ -181,9 +181,9 @@ def body_fn( def unbalanced_dykstra_kernel( - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, + k_q: jax.Array, + k_r: jax.Array, + k_g: jax.Array, gamma: float, ot_prob: linear_problem.LinearProblem, translation_invariant: bool = True, @@ -191,7 +191,7 @@ def unbalanced_dykstra_kernel( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array, jax.Array]: """Dykstra's algorithm for the unbalanced :class:`~ott.solvers.linear.sinkhorn_lr.LRSinkhorn` in kernel mode. @@ -317,9 +317,9 @@ def body_fn( def compute_lambdas( - const: Constants, state: State, gamma: float, g: jnp.ndarray, *, + const: Constants, state: State, gamma: float, g: jax.Array, *, lse_mode: bool -) -> Tuple[float, float]: +) -> tuple[float, float]: """TODO.""" gamma_inv = 1.0 / gamma rho_a = const.rho_a diff --git a/src/ott/solvers/linear/semidiscrete.py b/src/ott/solvers/linear/semidiscrete.py index 644dca146..3708fe41c 100644 --- a/src/ott/solvers/linear/semidiscrete.py +++ b/src/ott/solvers/linear/semidiscrete.py @@ -13,7 +13,8 @@ # limitations under the License. import dataclasses import math -from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Union +from collections.abc import Callable +from typing import TYPE_CHECKING, Any import jax import jax.experimental.sparse as jesp @@ -87,8 +88,8 @@ class HardAssignmentOutput: """ ot_prob: linear_problem.LinearProblem paired_indices: jax.Array - f: Optional[jax.Array] = None - g: Optional[jax.Array] = None + f: jax.Array | None = None + g: jax.Array | None = None @property def matrix(self) -> jesp.BCOO: @@ -138,18 +139,18 @@ class SemidiscreteOutput: """ g: jax.Array prob: sdlp.SemidiscreteLinearProblem - it: Optional[int] = None - losses: Optional[jax.Array] = None - errors: Optional[jax.Array] = None - converged: Optional[bool] = None + it: int | None = None + losses: jax.Array | None = None + errors: jax.Array | None = None + converged: bool | None = None def sample( self, rng: jax.Array, num_samples: int, *, - epsilon: Optional[float] = None, - ) -> Union[sinkhorn.SinkhornOutput, HardAssignmentOutput]: + epsilon: float | None = None, + ) -> sinkhorn.SinkhornOutput | HardAssignmentOutput: """Sample a point cloud and compute the OT solution. Args: @@ -190,7 +191,7 @@ def sample( ) def to_dual_potentials( - self, epsilon: Optional[float] = None + self, epsilon: float | None = None ) -> potentials.DualPotentials: """Compute the dual potential function :math:`f`. @@ -303,19 +304,19 @@ class SemidiscreteSolver: batch_size: int optimizer: optax.GradientTransformation error_eval_every: int = 1000 - error_batch_size: Optional[int] = None + error_batch_size: int | None = None error_num_repeats: int = 16 threshold: float = 1e-3 potential_ema: float = 0.99 epsilon_scheduler: Callable[[jax.Array, jax.Array], jax.Array] = constant_epsilon_scheduler - callback: Optional[Callable[[SemidiscreteState], None]] = None + callback: Callable[[SemidiscreteState], None] | None = None def __call__( self, rng: jax.Array, prob: sdlp.SemidiscreteLinearProblem, - g_init: Optional[jax.Array] = None, + g_init: jax.Array | None = None, ) -> SemidiscreteOutput: """Run the semidiscrete solver. @@ -409,7 +410,7 @@ def step( prob: sdlp.SemidiscreteLinearProblem, *, compute_error: bool = False, - rng_error: Optional[jax.Array] = None, + rng_error: jax.Array | None = None, ) -> SemidiscreteState: """Perform one optimization step. @@ -497,7 +498,7 @@ def _semidiscrete_loss( def _semidiscrete_loss_fwd( g: jax.Array, prob: linear_problem.LinearProblem, -) -> Tuple[jax.Array, Tuple[jax.Array, linear_problem.LinearProblem]]: +) -> tuple[jax.Array, tuple[jax.Array, linear_problem.LinearProblem]]: f, z = prob._c_transform(g, axis=1) # we assume uniform weights for `prob.a` return -(jnp.mean(f) + jnp.dot(g, prob.b)), (z, prob) @@ -506,7 +507,7 @@ def _semidiscrete_loss_fwd( def _semidiscrete_loss_bwd( res: jax.Array, g: jax.Array, -) -> Tuple[jax.Array, None]: +) -> tuple[jax.Array, None]: def soft_grad(z: jax.Array) -> jax.Array: if prob._b is None: # uniform weights @@ -553,7 +554,7 @@ def _marginal_chi2_error( batch_size: int, ) -> jax.Array: - def compute_chi2(matrix: Union[jax.Array, jesp.BCOO]) -> jax.Array: + def compute_chi2(matrix: jax.Array | jesp.BCOO) -> jax.Array: """Compute chi2 metric. Implements Eq. 3.5 in https://arxiv.org/pdf/2509.25519v1, @@ -577,7 +578,7 @@ def compute_chi2(matrix: Union[jax.Array, jesp.BCOO]) -> jax.Array: out = jnp.sum(out / prob.b) / (batch_size * (batch_size - 1.0)) return out - 1.0 - def body(chi2_err_avg: jax.Array, it: jax.Array) -> Tuple[jax.Array, None]: + def body(chi2_err_avg: jax.Array, it: jax.Array) -> tuple[jax.Array, None]: rng_it = jr.fold_in(rng, it) matrix = out.sample(rng_it, batch_size).matrix chi2 = compute_chi2(matrix) diff --git a/src/ott/solvers/linear/sinkhorn.py b/src/ott/solvers/linear/sinkhorn.py index 37f45d342..9d93a719c 100644 --- a/src/ott/solvers/linear/sinkhorn.py +++ b/src/ott/solvers/linear/sinkhorn.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, NamedTuple, Optional, Sequence, Tuple +from collections.abc import Callable, Sequence +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -30,16 +31,16 @@ __all__ = ["Sinkhorn", "SinkhornOutput"] ProgressFunction = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "SinkhornState"]], None] + [tuple[np.ndarray, np.ndarray, np.ndarray, "SinkhornState"]], None] class SinkhornState(NamedTuple): """Holds the state variables used to solve OT with Sinkhorn.""" - potentials: Tuple[jnp.ndarray, ...] - errors: Optional[jnp.ndarray] = None - old_fus: Optional[jnp.ndarray] = None - old_mapped_fus: Optional[jnp.ndarray] = None + potentials: tuple[jax.Array, ...] + errors: jax.Array | None = None + old_fus: jax.Array | None = None + old_mapped_fus: jax.Array | None = None def set(self, **kwargs: Any) -> "SinkhornState": """Return a copy of self, with potential overwrites.""" @@ -53,7 +54,7 @@ def solution_error( lse_mode: bool, parallel_dual_updates: bool, recenter: bool, - ) -> jnp.ndarray: + ) -> jax.Array: """State dependent function to return error.""" fu, gv = self.fu, self.gv if recenter and lse_mode: @@ -70,15 +71,15 @@ def solution_error( def compute_kl_reg_cost( # noqa: D102 self, ot_prob: linear_problem.LinearProblem, lse_mode: bool - ) -> jnp.ndarray: + ) -> jax.Array: return compute_kl_reg_cost(self.fu, self.gv, ot_prob, lse_mode) def recenter( self, - f: jnp.ndarray, - g: jnp.ndarray, + f: jax.Array, + g: jax.Array, ot_prob: linear_problem.LinearProblem, - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array]: """Re-center dual potentials. If the ``ot_prob`` is balanced, the ``f`` potential is zero-centered. @@ -114,25 +115,25 @@ def recenter( return f + shift, g - shift @property - def fu(self) -> jnp.ndarray: + def fu(self) -> jax.Array: """The first dual potential or scaling.""" return self.potentials[0] @property - def gv(self) -> jnp.ndarray: + def gv(self) -> jax.Array: """The second dual potential or scaling.""" return self.potentials[1] def solution_error( - f_u: jnp.ndarray, - g_v: jnp.ndarray, + f_u: jax.Array, + g_v: jax.Array, ot_prob: linear_problem.LinearProblem, *, norm_error: Sequence[int], lse_mode: bool, parallel_dual_updates: bool, -) -> jnp.ndarray: +) -> jax.Array: """Given two potential/scaling solutions, computes deviation to optimality. When the ``ot_prob`` problem is balanced and the usual Sinkhorn updates are @@ -146,8 +147,8 @@ def solution_error( additional quantities to qualify optimality must be taken into account. Args: - f_u: jnp.ndarray, potential or scaling - g_v: jnp.ndarray, potential or scaling + f_u: jax.Array, potential or scaling + g_v: jax.Array, potential or scaling ot_prob: linear OT problem norm_error: int, p-norm used to compute error. lse_mode: True if log-sum-exp operations, False if kernel vector products. @@ -189,9 +190,9 @@ def solution_error( def marginal_error( - f_u: jnp.ndarray, - g_v: jnp.ndarray, - target: jnp.ndarray, + f_u: jax.Array, + g_v: jax.Array, + target: jax.Array, geom: geometry.Geometry, axis: int = 0, norm_error: Sequence[int] = (1,), @@ -222,9 +223,9 @@ def marginal_error( def compute_kl_reg_cost( - f: jnp.ndarray, g: jnp.ndarray, ot_prob: linear_problem.LinearProblem, + f: jax.Array, g: jax.Array, ot_prob: linear_problem.LinearProblem, lse_mode: bool -) -> jnp.ndarray: +) -> jax.Array: r"""Compute objective of Sinkhorn for OT problem given dual solutions. The objective is evaluated for dual solution ``f`` and ``g``, using @@ -236,8 +237,8 @@ def compute_kl_reg_cost( values, ``jnp.where`` is used to cancel these contributions. Args: - f: jnp.ndarray, potential - g: jnp.ndarray, potential + f: jax.Array, potential + g: jax.Array, potential ot_prob: linear optimal transport problem. lse_mode: bool, whether to compute total mass in lse or kernel mode. @@ -314,13 +315,13 @@ class SinkhornOutput(NamedTuple): computations of errors. """ - potentials: Tuple[jnp.ndarray, ...] - errors: Optional[jnp.ndarray] = None - reg_ot_cost: Optional[jnp.ndarray] = None - ot_prob: Optional[linear_problem.LinearProblem] = None - threshold: Optional[jnp.ndarray] = None - converged: Optional[bool] = None - inner_iterations: Optional[int] = None + potentials: tuple[jax.Array, ...] + errors: jax.Array | None = None + reg_ot_cost: jax.Array | None = None + ot_prob: linear_problem.LinearProblem | None = None + threshold: jax.Array | None = None + converged: bool | None = None + inner_iterations: int | None = None def set(self, **kwargs: Any) -> "SinkhornOutput": """Return a copy of self, with potential overwrites.""" @@ -335,7 +336,7 @@ def set_cost( # noqa: D102 return self.set(reg_ot_cost=compute_kl_reg_cost(f, g, ot_prob, lse_mode)) @property - def dual_cost(self) -> jnp.ndarray: + def dual_cost(self) -> jax.Array: """Return dual transport cost, without considering regularizer.""" a, b = self.ot_prob.a, self.ot_prob.b dual_cost = jnp.sum(jnp.where(a > 0.0, a * self.f, 0)) @@ -343,12 +344,12 @@ def dual_cost(self) -> jnp.ndarray: return dual_cost @property - def primal_cost(self) -> jnp.ndarray: + def primal_cost(self) -> jax.Array: """Return transport cost of current transport solution at geometry.""" return self.transport_cost_at_geom(other_geom=self.geom) @property - def ent_reg_cost(self) -> jnp.ndarray: + def ent_reg_cost(self) -> jax.Array: r"""Entropy regularized cost. This outputs @@ -370,7 +371,7 @@ def ent_reg_cost(self) -> jnp.ndarray: return self.reg_ot_cost - self.geom.epsilon * (ent_a + ent_b) @property - def kl_reg_cost(self) -> jnp.ndarray: + def kl_reg_cost(self) -> jax.Array: r"""KL regularized OT transport cost. This outputs @@ -392,9 +393,7 @@ def kl_reg_cost(self) -> jnp.ndarray: """ return self.reg_ot_cost - def transport_cost_at_geom( - self, other_geom: geometry.Geometry - ) -> jnp.ndarray: + def transport_cost_at_geom(self, other_geom: geometry.Geometry) -> jax.Array: r"""Return bare transport cost of current solution at any geometry. In order to compute cost, we check first if the geometry can be converted @@ -421,11 +420,11 @@ def geom(self) -> geometry.Geometry: # noqa: D102 return self.ot_prob.geom @property - def a(self) -> jnp.ndarray: # noqa: D102 + def a(self) -> jax.Array: # noqa: D102 return self.ot_prob.a @property - def b(self) -> jnp.ndarray: # noqa: D102 + def b(self) -> jax.Array: # noqa: D102 return self.ot_prob.b @property @@ -434,13 +433,13 @@ def n_iters(self) -> int: # noqa: D102 return jnp.sum(self.errors != -1) * self.inner_iterations @property - def scalings(self) -> Tuple[jnp.ndarray, jnp.ndarray]: # noqa: D102 + def scalings(self) -> tuple[jax.Array, jax.Array]: # noqa: D102 u = self.ot_prob.geom.scaling_from_potential(self.f) v = self.ot_prob.geom.scaling_from_potential(self.g) return u, v @property - def matrix(self) -> jnp.ndarray: + def matrix(self) -> jax.Array: """Transport matrix if it can be instantiated.""" try: return self.ot_prob.geom.transport_from_potentials(self.f, self.g) @@ -448,16 +447,16 @@ def matrix(self) -> jnp.ndarray: return self.ot_prob.geom.transport_from_scalings(*self.scalings) @property - def transport_mass(self) -> jnp.ndarray: + def transport_mass(self) -> jax.Array: """Sum of transport matrix.""" return self.marginal(0).sum() def apply( self, - inputs: jnp.ndarray, + inputs: jax.Array, axis: int = 0, lse_mode: bool = True - ) -> jnp.ndarray: + ) -> jax.Array: """Apply the transport to a ndarray; axis=1 for its transpose.""" geom = self.ot_prob.geom if lse_mode: @@ -468,10 +467,10 @@ def apply( v = geom.scaling_from_potential(self.g) return geom.apply_transport_from_scalings(u, v, inputs, axis=axis) - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + def marginal(self, axis: int) -> jax.Array: # noqa: D102 return self.ot_prob.geom.marginal_from_potentials(self.f, self.g, axis=axis) - def cost_at_geom(self, other_geom: geometry.Geometry) -> jnp.ndarray: + def cost_at_geom(self, other_geom: geometry.Geometry) -> jax.Array: """Return reg-OT cost for matrix, evaluated at other cost matrix.""" return ( jnp.sum(self.matrix * other_geom.cost_matrix) - @@ -479,7 +478,7 @@ def cost_at_geom(self, other_geom: geometry.Geometry) -> jnp.ndarray: ) def to_dual_potentials( - self, epsilon: Optional[float] = None + self, epsilon: float | None = None ) -> potentials.DualPotentials: """Compute dual potential functions. @@ -500,17 +499,17 @@ def to_dual_potentials( return potentials.DualPotentials(f_fn, g_fn, cost_fn=cost_fn) @property - def f(self) -> jnp.ndarray: + def f(self) -> jax.Array: """The first dual potential.""" return self.potentials[0] @property - def g(self) -> jnp.ndarray: + def g(self) -> jax.Array: """The second dual potential.""" return self.potentials[1] @property - def entropy(self) -> jnp.ndarray: + def entropy(self) -> jax.Array: """Entropy of the coupling.""" marginal_a = self.marginal(1) marginal_b = self.marginal(0) @@ -522,14 +521,14 @@ def entropy(self) -> jnp.ndarray: ) / self.geom.epsilon @property - def normalized_entropy(self) -> jnp.ndarray: + def normalized_entropy(self) -> jax.Array: """Renormalized entropy of coupling when the problem is assignment.""" is_assign = self.ot_prob.is_assignment assert is_assign, "Normalized entropy only valid for assignment problem." return self.entropy / jnp.log(self.geom.shape[0]) - 1.0 @property - def diag(self) -> jnp.ndarray: + def diag(self) -> jax.Array: """Diagonal of the transport matrix.""" assert self.ot_prob.geom.is_square, ( "Problem must be square for ", "transport matrix to have a diag." @@ -723,15 +722,15 @@ def __init__( inner_iterations: int = 10, min_iterations: int = 0, max_iterations: int = 2000, - momentum: Optional[acceleration.Momentum] = None, - anderson: Optional[acceleration.AndersonAcceleration] = None, + momentum: acceleration.Momentum | None = None, + anderson: acceleration.AndersonAcceleration | None = None, parallel_dual_updates: bool = False, recenter_potentials: bool = False, - use_danskin: Optional[bool] = None, - implicit_diff: Optional[implicit_lib.ImplicitDiff - ] = implicit_lib.ImplicitDiff(), # noqa: B008 - initializer: Optional[init_lib.SinkhornInitializer] = None, - progress_fn: Optional[ProgressFunction] = None, + use_danskin: bool | None = None, + implicit_diff: implicit_lib.ImplicitDiff + | None = implicit_lib.ImplicitDiff(), # noqa: B008 + initializer: init_lib.SinkhornInitializer | None = None, + progress_fn: ProgressFunction | None = None, ): self.lse_mode = lse_mode self.threshold = threshold @@ -782,7 +781,7 @@ def __init__( def __call__( self, ot_prob: linear_problem.LinearProblem, - init: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None, + init: tuple[jax.Array, jax.Array] | None = None, **kwargs: Any, ) -> SinkhornOutput: """Run Sinkhorn algorithm. @@ -815,9 +814,7 @@ def xi(tau_i: float, tau_j: float) -> float: k_ij = k(tau_i, tau_j) return k_ij / (1.0 - k_ij) - def smin( - potential: jnp.ndarray, marginal: jnp.ndarray, tau: float - ) -> float: + def smin(potential: jax.Array, marginal: jax.Array, tau: float) -> float: rho = uf.rho(ot_prob.epsilon, tau) return -rho * mu.logsumexp(-potential / rho, b=marginal) @@ -962,8 +959,8 @@ def outer_iterations(self) -> int: return np.ceil(self.max_iterations / self.inner_iterations).astype(int) def init_state( - self, ot_prob: linear_problem.LinearProblem, init: Tuple[jnp.ndarray, - jnp.ndarray] + self, ot_prob: linear_problem.LinearProblem, init: tuple[jax.Array, + jax.Array] ) -> SinkhornState: """Return the initial state of the loop.""" errors = -jnp.ones((self.outer_iterations, len(self.norm_error)), @@ -1029,7 +1026,7 @@ def output_from_state( inner_iterations=self.inner_iterations) @property - def norm_error(self) -> Tuple[int, ...]: + def norm_error(self) -> tuple[int, ...]: """Powers used to compute the p-norm between marginal/target.""" # To change momentum adaptively, one needs errors in ||.||_1 norm. # In that case, we add this exponent to the list of errors to compute, @@ -1051,7 +1048,7 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 def run( ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] + init: tuple[jax.Array, ...] ) -> SinkhornOutput: """Run loop of the solver, outputting a state upgraded to an output.""" iter_fun = _iterations_implicit if solver.implicit_diff else iterations @@ -1064,19 +1061,19 @@ def run( def iterations( ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] + init: tuple[jax.Array, ...] ) -> SinkhornOutput: """Jittable Sinkhorn loop. args contain initialization variables.""" def cond_fn( - iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], + iteration: int, const: tuple[linear_problem.LinearProblem, Sinkhorn], state: SinkhornState ) -> bool: _, solver = const return solver._continue(state, iteration) def body_fn( - iteration: int, const: Tuple[linear_problem.LinearProblem, Sinkhorn], + iteration: int, const: tuple[linear_problem.LinearProblem, Sinkhorn], state: SinkhornState, compute_error: bool ) -> SinkhornState: ot_prob, solver = const @@ -1101,8 +1098,8 @@ def body_fn( def _iterations_taped( ot_prob: linear_problem.LinearProblem, solver: Sinkhorn, - init: Tuple[jnp.ndarray, ...] -) -> Tuple[SinkhornOutput, Tuple[jnp.ndarray, jnp.ndarray, + init: tuple[jax.Array, ...] +) -> tuple[SinkhornOutput, tuple[jax.Array, jax.Array, linear_problem.LinearProblem, Sinkhorn]]: """Run forward pass of the Sinkhorn algorithm storing side information.""" state = iterations(ot_prob, solver, init) @@ -1121,7 +1118,7 @@ def _iterations_implicit_bwd(res, gr: SinkhornOutput): considered. Returns: - a tuple of gradients: PyTree for geom, one jnp.ndarray for each of a and b. + a tuple of gradients: PyTree for geom, one jax.Array for each of a and b. """ f, g, ot_prob, solver = res out = solver.implicit_diff.gradient( diff --git a/src/ott/solvers/linear/sinkhorn_lr.py b/src/ott/solvers/linear/sinkhorn_lr.py index 7cadf826a..1abfa90cd 100644 --- a/src/ott/solvers/linear/sinkhorn_lr.py +++ b/src/ott/solvers/linear/sinkhorn_lr.py @@ -11,7 +11,8 @@ # 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 typing import Any, Callable, Mapping, NamedTuple, Optional, Tuple +from collections.abc import Callable, Mapping +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -28,17 +29,17 @@ __all__ = ["LRSinkhorn", "LRSinkhornOutput"] ProgressFunction = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRSinkhornState"]], None] + [tuple[np.ndarray, np.ndarray, np.ndarray, "LRSinkhornState"]], None] class LRSinkhornState(NamedTuple): """State of the Low Rank Sinkhorn algorithm.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray + q: jax.Array + r: jax.Array + g: jax.Array gamma: float - costs: jnp.ndarray - errors: jnp.ndarray + costs: jax.Array + errors: jax.Array crossed_threshold: bool def compute_error( # noqa: D102 @@ -69,8 +70,8 @@ def reg_ot_cost( # noqa: D102 ) def solution_error( # noqa: D102 - self, ot_prob: linear_problem.LinearProblem, norm_error: Tuple[int, ...] - ) -> jnp.ndarray: + self, ot_prob: linear_problem.LinearProblem, norm_error: tuple[int, ...] + ) -> jax.Array: return solution_error(self.q, self.r, ot_prob, norm_error) def set(self, **kwargs: Any) -> "LRSinkhornState": @@ -79,9 +80,9 @@ def set(self, **kwargs: Any) -> "LRSinkhornState": def compute_reg_ot_cost( - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, + q: jax.Array, + r: jax.Array, + g: jax.Array, ot_prob: linear_problem.LinearProblem, epsilon: float, use_danskin: bool = True @@ -118,9 +119,9 @@ def compute_reg_ot_cost( def solution_error( - q: jnp.ndarray, r: jnp.ndarray, ot_prob: linear_problem.LinearProblem, - norm_error: Tuple[int, ...] -) -> jnp.ndarray: + q: jax.Array, r: jax.Array, ot_prob: linear_problem.LinearProblem, + norm_error: tuple[int, ...] +) -> jax.Array: """Compute solution error. Since only balanced case is available for LR, this is marginal deviation. @@ -153,19 +154,19 @@ def solution_error( class LRSinkhornOutput(NamedTuple): """Transport interface for a low-rank Sinkhorn solution.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - costs: jnp.ndarray + q: jax.Array + r: jax.Array + g: jax.Array + costs: jax.Array # TODO(michalk8): must be called `errors`, because of `store_inner_errors` # in future, enforce via class hierarchy - errors: jnp.ndarray + errors: jax.Array ot_prob: linear_problem.LinearProblem epsilon: float inner_iterations: int converged: bool # TODO(michalk8): Optional is an artifact of the current impl., refactor - reg_ot_cost: Optional[float] = None + reg_ot_cost: float | None = None def set(self, **kwargs: Any) -> "LRSinkhornOutput": """Return a copy of self, with potential overwrites.""" @@ -199,11 +200,11 @@ def geom(self) -> geometry.Geometry: # noqa: D102 return self.ot_prob.geom @property - def a(self) -> jnp.ndarray: # noqa: D102 + def a(self) -> jax.Array: # noqa: D102 return self.ot_prob.a @property - def b(self) -> jnp.ndarray: # noqa: D102 + def b(self) -> jax.Array: # noqa: D102 return self.ot_prob.b @property @@ -211,17 +212,17 @@ def n_iters(self) -> int: # noqa: D102 return jnp.sum(self.errors != -1) * self.inner_iterations @property - def matrix(self) -> jnp.ndarray: + def matrix(self) -> jax.Array: """Transport matrix if it can be instantiated.""" return (self.q * self._inv_g) @ self.r.T - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + def apply(self, inputs: jax.Array, axis: int = 0) -> jax.Array: """Apply the transport to a array; axis=1 for its transpose.""" q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) # for `axis=0`: (batch, m), (m, r), (r,), (r, n) return ((inputs @ r) * self._inv_g) @ q.T - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + def marginal(self, axis: int) -> jax.Array: # noqa: D102 length = self.q.shape[0] if axis == 0 else self.r.shape[0] return self.apply(jnp.ones(length,), axis=axis) @@ -244,7 +245,7 @@ def transport_mass(self) -> float: return self.marginal(0).sum() @property - def _inv_g(self) -> jnp.ndarray: + def _inv_g(self) -> jax.Array: return 1.0 / self.g @@ -296,12 +297,12 @@ def __init__( gamma: float = 10.0, gamma_rescale: bool = True, epsilon: float = 0.0, - initializer: Optional[initializers_lr.LRInitializer] = None, + initializer: initializers_lr.LRInitializer | None = None, lse_mode: bool = True, inner_iterations: int = 10, use_danskin: bool = True, - kwargs_dys: Optional[Mapping[str, Any]] = None, - progress_fn: Optional[ProgressFunction] = None, + kwargs_dys: Mapping[str, Any] | None = None, + progress_fn: ProgressFunction | None = None, **kwargs: Any, ): kwargs["implicit_diff"] = None # not yet implemented @@ -324,7 +325,7 @@ def __init__( def __call__( self, ot_prob: linear_problem.LinearProblem, - init: Optional[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]] = None, + init: tuple[jax.Array, jax.Array, jax.Array] | None = None, **kwargs: Any, ) -> LRSinkhornOutput: """Run low-rank Sinkhorn. @@ -351,7 +352,7 @@ def _get_costs( self, ot_prob: linear_problem.LinearProblem, state: LRSinkhornState, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: + ) -> tuple[jax.Array, jax.Array, jax.Array, float]: log_q, log_r, log_g = ( mu.safe_log(state.q), mu.safe_log(state.r), mu.safe_log(state.g) ) @@ -387,9 +388,9 @@ def _get_costs( # TODO(michalk8): move to `lr_utils` when refactoring this def dykstra_update_lse( self, - c_q: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, + c_q: jax.Array, + c_r: jax.Array, + h: jax.Array, gamma: float, ot_prob: linear_problem.LinearProblem, min_entry_value: float = 1e-6, @@ -397,7 +398,7 @@ def dykstra_update_lse( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Run Dykstra's algorithm.""" # shortcuts for problem's definition. r = self.rank @@ -415,24 +416,24 @@ def dykstra_update_lse( constants = c_q, c_r, loga, logb def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...] ) -> bool: del iteration, constants *_, err = state_inner return err > tolerance def _softm( - f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int - ) -> jnp.ndarray: + f: jax.Array, g: jax.Array, c: jax.Array, axis: int + ) -> jax.Array: return jsp.special.logsumexp( gamma * (f[:, None] + g[None, :] - c), axis=axis ) def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...], compute_error: bool + ) -> tuple[jax.Array, ...]: # TODO(michalk8): in the future, use `NamedTuple` f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner c_q, c_r, loga, logb = constants @@ -481,15 +482,15 @@ def body_fn( return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err def recompute_couplings( - f1: jnp.ndarray, - g1: jnp.ndarray, - c_q: jnp.ndarray, - f2: jnp.ndarray, - g2: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, + f1: jax.Array, + g1: jax.Array, + c_q: jax.Array, + f2: jax.Array, + g2: jax.Array, + c_r: jax.Array, + h: jax.Array, gamma: float, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) g = jnp.exp(gamma * h) @@ -504,9 +505,9 @@ def recompute_couplings( def dykstra_update_kernel( self, - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, + k_q: jax.Array, + k_r: jax.Array, + k_g: jax.Array, gamma: float, ot_prob: linear_problem.LinearProblem, min_entry_value: float = 1e-6, @@ -514,7 +515,7 @@ def dykstra_update_kernel( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Run Dykstra's algorithm.""" # shortcuts for problem's definition. rank = self.rank @@ -533,17 +534,17 @@ def dykstra_update_kernel( constants = k_q, k_r, k_g, a, b def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...] ) -> bool: del iteration, constants *_, err = state_inner return err > tolerance def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...], compute_error: bool + ) -> tuple[jax.Array, ...]: # TODO(michalk8): in the future, use `NamedTuple` u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner k_q, k_r, k_g, a, b = constants @@ -580,14 +581,14 @@ def body_fn( return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err def recompute_couplings( - u1: jnp.ndarray, - v1: jnp.ndarray, - k_q: jnp.ndarray, - u2: jnp.ndarray, - v2: jnp.ndarray, - k_r: jnp.ndarray, - g: jnp.ndarray, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + u1: jax.Array, + v1: jax.Array, + k_q: jax.Array, + u2: jax.Array, + v2: jax.Array, + k_r: jax.Array, + g: jax.Array, + ) -> tuple[jax.Array, jax.Array, jax.Array]: q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) return q, r, g @@ -694,12 +695,12 @@ def one_iteration( return state @property - def norm_error(self) -> Tuple[int]: # noqa: D102 + def norm_error(self) -> tuple[int]: # noqa: D102 return self._norm_error, def init_state( self, ot_prob: linear_problem.LinearProblem, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] + init: tuple[jax.Array, jax.Array, jax.Array] ) -> LRSinkhornState: """Return the initial state of the loop.""" q, r, g = init @@ -779,7 +780,7 @@ def _diverged(self, state: LRSinkhornState, iteration: int) -> bool: def run( ot_prob: linear_problem.LinearProblem, solver: LRSinkhorn, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + init: tuple[jax.Array, jax.Array, jax.Array], ) -> LRSinkhornOutput: """Run loop of the solver, outputting a state upgraded to an output.""" out = sinkhorn.iterations(ot_prob, solver, init) diff --git a/src/ott/solvers/linear/univariate.py b/src/ott/solvers/linear/univariate.py index 9f26d4e76..890e8627b 100644 --- a/src/ott/solvers/linear/univariate.py +++ b/src/ott/solvers/linear/univariate.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import NamedTuple, Optional, Tuple +from typing import NamedTuple import jax import jax.experimental.sparse as jesp @@ -54,11 +54,11 @@ class UnivariateOutput(NamedTuple): dual_b: Array of shape ``[m,]`` containing the second dual variable. """ prob: linear_problem.LinearProblem - ot_costs: jnp.ndarray - paired_indices: Optional[jnp.ndarray] = None - mass_paired_indices: Optional[jnp.ndarray] = None - dual_a: Optional[jnp.ndarray] = None - dual_b: Optional[jnp.ndarray] = None + ot_costs: jax.Array + paired_indices: jax.Array | None = None + mass_paired_indices: jax.Array | None = None + dual_a: jax.Array | None = None + dual_b: jax.Array | None = None @property def transport_matrices(self) -> jesp.BCOO: @@ -80,7 +80,7 @@ def mean_transport_matrix(self) -> jesp.BCOO: return sparse_mean(self.transport_matrices, axis=0) @property - def dual_costs(self) -> jnp.ndarray: + def dual_costs(self) -> jax.Array: """Array of shape ``[d,]`` containing the dual costs.""" assert self.dual_a is not None, "Dual variables have not been computed." dual_obj = jnp.sum(self.dual_a * self.prob.a[None, :], axis=1) @@ -112,7 +112,7 @@ def uniform_solver( @functools.partial(jax.vmap, in_axes=[1, 1]) @functools.partial(jax.vmap, in_axes=[0, 0]) - def cost(x: jnp.ndarray, y: jnp.ndarray) -> float: + def cost(x: jax.Array, y: jax.Array) -> float: return cost_fn(x[None], y[None]) assert prob.is_equal_size, "Source and target have different sizes." @@ -168,7 +168,7 @@ def quantile_solver( """ # noqa: E501 @functools.partial(jax.vmap, in_axes=[1, 1]) - def dist(x: jnp.ndarray, y: jnp.ndarray): + def dist(x: jax.Array, y: jax.Array): x, i_x = mu.sort_and_argsort(x, argsort=True) y, i_y = mu.sort_and_argsort(y, argsort=True) @@ -241,24 +241,24 @@ def north_west_solver(prob: linear_problem.LinearProblem) -> UnivariateOutput: """ # noqa: E501 class State(NamedTuple): - x: jnp.ndarray - y: jnp.ndarray - a: jnp.ndarray - b: jnp.ndarray - paired_indices: jnp.ndarray - mass_paired_indices: jnp.ndarray - dual_a: jnp.ndarray - dual_b: jnp.ndarray + x: jax.Array + y: jax.Array + a: jax.Array + b: jax.Array + paired_indices: jax.Array + mass_paired_indices: jax.Array + dual_a: jax.Array + dual_b: jax.Array def dual_a_update(state: State, i: int, - j: int) -> Tuple[State, jnp.ndarray, jnp.ndarray]: + j: int) -> tuple[State, jax.Array, jax.Array]: next_ixs = jnp.array([i + 1, j]) val = cost_fn(state.x[i + 1, None], state.y[j, None]) - state.dual_b[j] da = state.dual_a.at[i + 1].set(val) return state._replace(dual_a=da), state.a[i], next_ixs def dual_b_update(state: State, i: int, - j: int) -> Tuple[State, jnp.ndarray, jnp.ndarray]: + j: int) -> tuple[State, jax.Array, jax.Array]: next_ixs = jnp.array([i, j + 1]) val = cost_fn(state.x[i, None], state.y[j + 1, None]) - state.dual_a[i] db = state.dual_b.at[j + 1].set(val) @@ -280,7 +280,7 @@ def body_fun(ix: int, state: State) -> State: ) @functools.partial(jax.vmap, in_axes=[1, 1]) - def dist(x: jnp.ndarray, y: jnp.ndarray): + def dist(x: jax.Array, y: jax.Array): x, i_x = mu.sort_and_argsort(x, argsort=True) y, i_y = mu.sort_and_argsort(y, argsort=True) sorted_a, sorted_b = a[i_x], b[i_y] diff --git a/src/ott/solvers/quadratic/_solve.py b/src/ott/solvers/quadratic/_solve.py index f12569f6f..f42d0f5ac 100644 --- a/src/ott/solvers/quadratic/_solve.py +++ b/src/ott/solvers/quadratic/_solve.py @@ -11,9 +11,9 @@ # 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 typing import Any, Dict, Literal, Optional, Union +from typing import Any, Literal -import jax.numpy as jnp +import jax from ott.geometry import geometry from ott.problems.quadratic import quadratic_costs, quadratic_problem @@ -27,18 +27,18 @@ def solve( geom_xx: geometry.Geometry, geom_yy: geometry.Geometry, - geom_xy: Optional[geometry.Geometry] = None, + geom_xy: geometry.Geometry | None = None, fused_penalty: float = 1.0, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, + b: jax.Array | None = None, tau_a: float = 1.0, tau_b: float = 1.0, - loss: Union[Literal["sqeucl", "kl"], quadratic_costs.GWLoss] = "sqeucl", + loss: Literal["sqeucl", "kl"] | quadratic_costs.GWLoss = "sqeucl", gw_unbalanced_correction: bool = True, rank: int = -1, - linear_solver_kwargs: Optional[Dict[str, Any]] = None, + linear_solver_kwargs: dict[str, Any] | None = None, **kwargs: Any, -) -> Union[gw.GWOutput, lrgw.LRGWOutput]: +) -> gw.GWOutput | lrgw.LRGWOutput: """Solve quadratic regularized OT problem using a Gromov-Wasserstein solver. Args: diff --git a/src/ott/solvers/quadratic/gromov_wasserstein.py b/src/ott/solvers/quadratic/gromov_wasserstein.py index a91de37c1..bf0bdbce1 100644 --- a/src/ott/solvers/quadratic/gromov_wasserstein.py +++ b/src/ott/solvers/quadratic/gromov_wasserstein.py @@ -11,17 +11,8 @@ # 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 typing import ( - Any, - Callable, - Dict, - Literal, - NamedTuple, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Sequence +from typing import Any, Literal, NamedTuple import jax import jax.numpy as jnp @@ -37,10 +28,10 @@ __all__ = ["GromovWasserstein", "GWOutput"] -LinearOutput = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput] +LinearOutput = sinkhorn.SinkhornOutput | sinkhorn_lr.LRSinkhornOutput ProgressCallbackFn = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "GWState"]], None] + [tuple[np.ndarray, np.ndarray, np.ndarray, "GWState"]], None] class GWOutput(NamedTuple): @@ -60,12 +51,12 @@ class GWOutput(NamedTuple): old_transport_mass: Holds total mass of transport at previous iteration. """ - costs: Optional[jnp.ndarray] = None - linear_convergence: Optional[jnp.ndarray] = None + costs: jax.Array | None = None + linear_convergence: jax.Array | None = None converged: bool = False - errors: Optional[jnp.ndarray] = None - linear_state: Optional[LinearOutput] = None - geom: Optional[geometry.Geometry] = None + errors: jax.Array | None = None + linear_state: LinearOutput | None = None + geom: geometry.Geometry | None = None # Intermediate values. old_transport_mass: float = 1.0 @@ -74,11 +65,11 @@ def set(self, **kwargs: Any) -> "GWOutput": return self._replace(**kwargs) @property - def matrix(self) -> jnp.ndarray: + def matrix(self) -> jax.Array: """Transport matrix.""" return self._rescale_factor * self.linear_state.matrix - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + def apply(self, inputs: jax.Array, axis: int = 0) -> jax.Array: """Apply the transport to an array; axis=1 for its transpose.""" return self._rescale_factor * self.linear_state.apply(inputs, axis=axis) @@ -117,12 +108,12 @@ class GWState(NamedTuple): at each iteration. """ - costs: jnp.ndarray - linear_convergence: jnp.ndarray + costs: jax.Array + linear_convergence: jax.Array linear_state: LinearOutput linear_pb: linear_problem.LinearProblem old_transport_mass: float - errors: Optional[jnp.ndarray] = None + errors: jax.Array | None = None def set(self, **kwargs: Any) -> "GWState": """Return a copy of self, possibly with overwrites.""" @@ -180,10 +171,10 @@ def __init__( self, linear_solver: sinkhorn.Sinkhorn, epsilon: float = 1.0, - relative_epsilon: Optional[Literal["mean", "std"]] = None, - initializer: Optional[quad_initializers.BaseQuadraticInitializer] = None, + relative_epsilon: Literal["mean", "std"] | None = None, + initializer: quad_initializers.BaseQuadraticInitializer | None = None, warm_start: bool = False, - progress_fn: Optional[ProgressCallbackFn] = None, + progress_fn: ProgressCallbackFn | None = None, **kwargs: Any ): super().__init__(linear_solver, **kwargs) @@ -197,7 +188,7 @@ def __init__( def __call__( self, prob: quadratic_problem.QuadraticProblem, - init: Optional[linear_problem.LinearProblem] = None, + init: linear_problem.LinearProblem | None = None, **kwargs: Any, ) -> GWOutput: """Run the Gromov-Wasserstein solver. @@ -293,7 +284,7 @@ def output_from_state( old_transport_mass=state.old_transport_mass ) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 children, aux_data = super().tree_flatten() aux_data["epsilon"] = self.epsilon aux_data["relative_epsilon"] = self.relative_epsilon @@ -304,7 +295,7 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "GromovWasserstein": linear_solver, threshold = children return cls(linear_solver, threshold=threshold, **aux_data) diff --git a/src/ott/solvers/quadratic/gromov_wasserstein_lr.py b/src/ott/solvers/quadratic/gromov_wasserstein_lr.py index 6a3c9db76..8fc9fdf92 100644 --- a/src/ott/solvers/quadratic/gromov_wasserstein_lr.py +++ b/src/ott/solvers/quadratic/gromov_wasserstein_lr.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """A Jax implementation of the unbalanced low-rank GW algorithm.""" -from typing import Any, Callable, Mapping, NamedTuple, Optional, Tuple +from collections.abc import Callable, Mapping +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -31,17 +32,17 @@ __all__ = ["LRGromovWasserstein", "LRGWOutput"] ProgressFunction = Callable[ - [Tuple[np.ndarray, np.ndarray, np.ndarray, "LRGWState"]], None] + [tuple[np.ndarray, np.ndarray, np.ndarray, "LRGWState"]], None] class LRGWState(NamedTuple): """State of the low-rank GW algorithm.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray + q: jax.Array + r: jax.Array + g: jax.Array gamma: float - costs: jnp.ndarray - errors: jnp.ndarray + costs: jax.Array + errors: jax.Array crossed_threshold: bool def compute_error( # noqa: D102 @@ -75,9 +76,9 @@ def set(self, **kwargs: Any) -> "LRGWState": def compute_reg_gw_cost( - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, + q: jax.Array, + r: jax.Array, + g: jax.Array, ot_prob: quadratic_problem.QuadraticProblem, epsilon: float, use_danskin: bool = False @@ -97,7 +98,7 @@ def compute_reg_gw_cost( regularized OT cost, the (primal) transport cost of the low-rank solution. """ - def ent(x: jnp.ndarray) -> float: + def ent(x: jax.Array) -> float: # generalized entropy return jnp.sum(jsp.special.entr(x) + x) @@ -130,18 +131,18 @@ def ent(x: jnp.ndarray) -> float: class LRGWOutput(NamedTuple): """Transport interface for a low-rank GW solution.""" - q: jnp.ndarray - r: jnp.ndarray - g: jnp.ndarray - costs: jnp.ndarray + q: jax.Array + r: jax.Array + g: jax.Array + costs: jax.Array # TODO(michalk8): must be called `errors`, because of `store_inner_errors` # in future, enforce via class hierarchy - errors: jnp.ndarray + errors: jax.Array ot_prob: quadratic_problem.QuadraticProblem epsilon: float inner_iterations: int converged: bool - reg_gw_cost: Optional[float] = None + reg_gw_cost: float | None = None def set(self, **kwargs: Any) -> "LRGWOutput": """Return a copy of self, with potential overwrites.""" @@ -176,11 +177,11 @@ def geom(self) -> geometry.Geometry: # noqa: D102 return _linearized_geometry(self.ot_prob, q=self.q, r=self.r, g=self.g) @property - def a(self) -> jnp.ndarray: # noqa: D102 + def a(self) -> jax.Array: # noqa: D102 return self.ot_prob.a @property - def b(self) -> jnp.ndarray: # noqa: D102 + def b(self) -> jax.Array: # noqa: D102 return self.ot_prob.b @property @@ -188,17 +189,17 @@ def n_iters(self) -> int: # noqa: D102 return jnp.sum(self.errors != -1) * self.inner_iterations @property - def matrix(self) -> jnp.ndarray: + def matrix(self) -> jax.Array: """Transport matrix if it can be instantiated.""" return (self.q * self._inv_g) @ self.r.T - def apply(self, inputs: jnp.ndarray, axis: int = 0) -> jnp.ndarray: + def apply(self, inputs: jax.Array, axis: int = 0) -> jax.Array: """Apply the transport to a array; axis=1 for its transpose.""" q, r = (self.q, self.r) if axis == 1 else (self.r, self.q) # for `axis=0`: (batch, m), (m, r), (r,), (r, n) return ((inputs @ r) * self._inv_g) @ q.T - def marginal(self, axis: int) -> jnp.ndarray: # noqa: D102 + def marginal(self, axis: int) -> jax.Array: # noqa: D102 length = self.q.shape[0] if axis == 0 else self.r.shape[0] return self.apply(jnp.ones(length,), axis=axis) @@ -236,7 +237,7 @@ def transport_mass(self) -> float: return self.marginal(0).sum() @property - def _inv_g(self) -> jnp.ndarray: + def _inv_g(self) -> jax.Array: return 1.0 / self.g @@ -289,15 +290,15 @@ def __init__( gamma: float = 10.0, gamma_rescale: bool = True, epsilon: float = 0.0, - initializer: Optional[initializers_lr.LRInitializer] = None, + initializer: initializers_lr.LRInitializer | None = None, lse_mode: bool = True, use_danskin: bool = True, implicit_diff: bool = False, inner_iterations: int = 2_000, min_iterations: int = 10_000, max_iterations: int = 100_000, - kwargs_dys: Optional[Mapping[str, Any]] = None, - progress_fn: Optional[ProgressFunction] = None, + kwargs_dys: Mapping[str, Any] | None = None, + progress_fn: ProgressFunction | None = None, **kwargs: Any, ): assert not implicit_diff, "Implicit diff. not yet implemented." @@ -324,8 +325,8 @@ def __init__( def __call__( self, ot_prob: quadratic_problem.QuadraticProblem, - init: Optional[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]] = None, - rng: Optional[jax.Array] = None, + init: tuple[jax.Array, jax.Array, jax.Array] | None = None, + rng: jax.Array | None = None, **kwargs: Any, ) -> LRGWOutput: """Run the low-rank Gromov-Wasserstein solver. @@ -359,7 +360,7 @@ def _get_costs( self, ot_prob: quadratic_problem.QuadraticProblem, state: LRGWState, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float]: + ) -> tuple[jax.Array, jax.Array, jax.Array, float]: q, r, g = state.q, state.r, state.g log_q, log_r, log_g = mu.safe_log(q), mu.safe_log(r), mu.safe_log(g) inv_g = 1.0 / g[None, :] @@ -420,9 +421,9 @@ def _get_costs( # TODO(michalk8): move to `lr_utils` when refactoring this the future def dykstra_update_lse( self, - c_q: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, + c_q: jax.Array, + c_r: jax.Array, + h: jax.Array, gamma: float, ot_prob: quadratic_problem.QuadraticProblem, min_entry_value: float = 1e-6, @@ -430,7 +431,7 @@ def dykstra_update_lse( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Run Dykstra's algorithm.""" # shortcuts for problem's definition. r = self.rank @@ -448,24 +449,24 @@ def dykstra_update_lse( constants = c_q, c_r, loga, logb def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...] ) -> bool: del iteration, constants *_, err = state_inner return err > tolerance def _softm( - f: jnp.ndarray, g: jnp.ndarray, c: jnp.ndarray, axis: int - ) -> jnp.ndarray: + f: jax.Array, g: jax.Array, c: jax.Array, axis: int + ) -> jax.Array: return jsp.special.logsumexp( gamma * (f[:, None] + g[None, :] - c), axis=axis ) def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...], compute_error: bool + ) -> tuple[jax.Array, ...]: # TODO(michalk8): in the future, use `NamedTuple` f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err = state_inner c_q, c_r, loga, logb = constants @@ -515,15 +516,15 @@ def body_fn( return f1, f2, g1_old, g2_old, h_old, w_gi, w_gp, w_q, w_r, err def recompute_couplings( - f1: jnp.ndarray, - g1: jnp.ndarray, - c_q: jnp.ndarray, - f2: jnp.ndarray, - g2: jnp.ndarray, - c_r: jnp.ndarray, - h: jnp.ndarray, + f1: jax.Array, + g1: jax.Array, + c_q: jax.Array, + f2: jax.Array, + g2: jax.Array, + c_r: jax.Array, + h: jax.Array, gamma: float, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: q = jnp.exp(gamma * (f1[:, None] + g1[None, :] - c_q)) r = jnp.exp(gamma * (f2[:, None] + g2[None, :] - c_r)) g = jnp.exp(gamma * h) @@ -538,9 +539,9 @@ def recompute_couplings( def dykstra_update_kernel( self, - k_q: jnp.ndarray, - k_r: jnp.ndarray, - k_g: jnp.ndarray, + k_q: jax.Array, + k_r: jax.Array, + k_g: jax.Array, gamma: float, ot_prob: quadratic_problem.QuadraticProblem, min_entry_value: float = 1e-6, @@ -548,7 +549,7 @@ def dykstra_update_kernel( min_iter: int = 0, inner_iter: int = 10, max_iter: int = 10000 - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array, jax.Array]: """Run Dykstra's algorithm.""" # shortcuts for problem's definition. del gamma @@ -568,17 +569,17 @@ def dykstra_update_kernel( constants = k_q, k_r, k_g, a, b def cond_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...] + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...] ) -> bool: del iteration, constants *_, err = state_inner return err > tolerance def body_fn( - iteration: int, constants: Tuple[jnp.ndarray, ...], - state_inner: Tuple[jnp.ndarray, ...], compute_error: bool - ) -> Tuple[jnp.ndarray, ...]: + iteration: int, constants: tuple[jax.Array, ...], + state_inner: tuple[jax.Array, ...], compute_error: bool + ) -> tuple[jax.Array, ...]: # TODO(michalk8): in the future, use `NamedTuple` u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err = state_inner k_q, k_r, k_g, a, b = constants @@ -616,14 +617,14 @@ def body_fn( return u1, u2, v1_old, v2_old, g_old, q_gi, q_gp, q_q, q_r, err def recompute_couplings( - u1: jnp.ndarray, - v1: jnp.ndarray, - k_q: jnp.ndarray, - u2: jnp.ndarray, - v2: jnp.ndarray, - k_r: jnp.ndarray, - g: jnp.ndarray, - ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + u1: jax.Array, + v1: jax.Array, + k_q: jax.Array, + u2: jax.Array, + v2: jax.Array, + k_r: jax.Array, + g: jax.Array, + ) -> tuple[jax.Array, jax.Array, jax.Array]: q = u1.reshape((-1, 1)) * k_q * v1.reshape((1, -1)) r = u2.reshape((-1, 1)) * k_r * v2.reshape((1, -1)) return q, r, g @@ -731,12 +732,12 @@ def one_iteration( return state @property - def norm_error(self) -> Tuple[int]: # noqa: D102 + def norm_error(self) -> tuple[int]: # noqa: D102 return self._norm_error, def init_state( self, ot_prob: quadratic_problem.QuadraticProblem, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] + init: tuple[jax.Array, jax.Array, jax.Array] ) -> LRGWState: """Return the initial state of the loop.""" q, r, g = init @@ -816,7 +817,7 @@ def _diverged(self, state: LRGWState, iteration: int) -> bool: def run( ot_prob: quadratic_problem.QuadraticProblem, solver: LRGromovWasserstein, - init: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], + init: tuple[jax.Array, jax.Array, jax.Array], ) -> LRGWOutput: """Run loop of the solver, outputting a state upgraded to an output.""" out = sinkhorn.iterations(ot_prob, solver, init) @@ -827,9 +828,9 @@ def run( def dykstra_solution_error( - q: jnp.ndarray, r: jnp.ndarray, ot_prob: quadratic_problem.QuadraticProblem, - norm_error: Tuple[int, ...] -) -> jnp.ndarray: + q: jax.Array, r: jax.Array, ot_prob: quadratic_problem.QuadraticProblem, + norm_error: tuple[int, ...] +) -> jax.Array: """Compute solution error. Since only balanced case is available for LR, this is marginal deviation. @@ -862,9 +863,9 @@ def dykstra_solution_error( def _linearized_geometry( prob: quadratic_problem.QuadraticProblem, *, - q: jnp.ndarray, - r: jnp.ndarray, - g: jnp.ndarray, + q: jax.Array, + r: jax.Array, + g: jax.Array, ) -> low_rank.LRCGeometry: inv_sqrt_g = 1.0 / jnp.sqrt(g[None, :]) diff --git a/src/ott/solvers/quadratic/gw_barycenter.py b/src/ott/solvers/quadratic/gw_barycenter.py index af0fb5325..1d1f87569 100644 --- a/src/ott/solvers/quadratic/gw_barycenter.py +++ b/src/ott/solvers/quadratic/gw_barycenter.py @@ -11,8 +11,9 @@ # 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 collections.abc import Sequence from functools import partial -from typing import Any, Dict, NamedTuple, Optional, Sequence, Tuple, Union +from typing import Any, NamedTuple import jax import jax.numpy as jnp @@ -45,13 +46,13 @@ class GWBarycenterState(NamedTuple): gw_convergence: Array of shape ``[max_iter,]`` containing the convergence of all GW problems at each iteration. """ - cost: Optional[jnp.ndarray] = None - x: Optional[jnp.ndarray] = None - a: Optional[jnp.ndarray] = None - errors: Optional[jnp.ndarray] = None - costs: Optional[jnp.ndarray] = None - costs_bary: Optional[jnp.ndarray] = None - gw_convergence: Optional[jnp.ndarray] = None + cost: jax.Array | None = None + x: jax.Array | None = None + a: jax.Array | None = None + errors: jax.Array | None = None + costs: jax.Array | None = None + costs_bary: jax.Array | None = None + gw_convergence: jax.Array | None = None def set(self, **kwargs: Any) -> "GWBarycenterState": """Return a copy of self, possibly with overwrites.""" @@ -116,10 +117,9 @@ def init_state( self, problem: gw_barycenter.GWBarycenterProblem, bar_size: int, - bar_init: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, - jnp.ndarray]]] = None, - a: Optional[jnp.ndarray] = None, - rng: Optional[jax.Array] = None, + bar_init: jax.Array | tuple[jax.Array, jax.Array] | None = None, + a: jax.Array | None = None, + rng: jax.Array | None = None, ) -> GWBarycenterState: """Initialize the (fused) Gromov-Wasserstein barycenter state. @@ -194,13 +194,13 @@ def update_state( iteration: int, problem: gw_barycenter.GWBarycenterProblem, store_errors: bool = True, - ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: + ) -> tuple[float, bool, jax.Array, jax.Array | None]: """Solve the (fused) Gromov-Wasserstein barycenter problem.""" def solve_gw( - state: GWBarycenterState, b: jnp.ndarray, y: jnp.ndarray, - f: Optional[jnp.ndarray] - ) -> Tuple[float, bool, jnp.ndarray, Optional[jnp.ndarray]]: + state: GWBarycenterState, b: jax.Array, y: jax.Array, + f: jax.Array | None + ) -> tuple[float, bool, jax.Array, jax.Array | None]: quad_problem = problem._create_problem(state, y=y, b=b, f=f) out = self.quadratic_solver(quad_problem) return ( @@ -247,7 +247,7 @@ def output_from_state(self, state: GWBarycenterState) -> GWBarycenterState: # will be refactored in the future to create an output return state - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([self.quadratic_solver, self.threshold], { "min_iterations": self.min_iterations, "max_iterations": self.max_iterations, @@ -257,9 +257,8 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @partial(jax.vmap, in_axes=[None, 0, None, 0, None]) def init_transports( - solver, rng: jax.Array, a: jnp.ndarray, b: jnp.ndarray, - epsilon: Optional[float] -) -> jnp.ndarray: + solver, rng: jax.Array, a: jax.Array, b: jax.Array, epsilon: float | None +) -> jax.Array: """Initialize random 2D point cloud and solve the linear OT problem. Args: @@ -293,7 +292,7 @@ def cond_fn( return solver._continue(state, iteration) def body_fn( - iteration, constants: Tuple[GromovWassersteinBarycenter, + iteration, constants: tuple[GromovWassersteinBarycenter, gw_barycenter.GWBarycenterProblem], state: GWBarycenterState, compute_error: bool ) -> GWBarycenterState: diff --git a/src/ott/solvers/quadratic/lower_bound.py b/src/ott/solvers/quadratic/lower_bound.py index 58ed6caa1..af2137dfc 100644 --- a/src/ott/solvers/quadratic/lower_bound.py +++ b/src/ott/solvers/quadratic/lower_bound.py @@ -11,7 +11,7 @@ # 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 typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any from ott.geometry import pointcloud from ott.problems.quadratic import quadratic_problem @@ -27,7 +27,7 @@ def third_lower_bound( prob: quadratic_problem.QuadraticProblem, distrib_cost: "distrib_costs.UnivariateWasserstein", - epsilon: Optional[float] = None, + epsilon: float | None = None, **kwargs: Any, ) -> sinkhorn.SinkhornOutput: """Computes the third lower bound distance from :cite:`memoli:11`, def. 6.3. diff --git a/src/ott/solvers/utils.py b/src/ott/solvers/utils.py index 276f65d27..11642709c 100644 --- a/src/ott/solvers/utils.py +++ b/src/ott/solvers/utils.py @@ -11,7 +11,7 @@ # 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 typing import Any, Literal, Optional, Tuple, Union +from typing import Any, Literal import jax import jax.numpy as jnp @@ -27,17 +27,17 @@ "uniform_sampler", ] -ScaleCost_t = Union[float, Literal["mean", "max_cost", "median"]] +ScaleCost_t = float | Literal["mean", "max_cost", "median"] def match_linear( - x: jnp.ndarray, - y: Optional[jnp.ndarray], - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, + x: jax.Array, + y: jax.Array | None, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, scale_cost: ScaleCost_t = 1.0, **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Compute solution to a linear OT problem. Args: @@ -59,14 +59,14 @@ def match_linear( def match_quadratic( - xx: jnp.ndarray, - yy: jnp.ndarray, - x: Optional[jnp.ndarray] = None, - y: Optional[jnp.ndarray] = None, + xx: jax.Array, + yy: jax.Array, + x: jax.Array | None = None, + y: jax.Array | None = None, scale_cost: ScaleCost_t = 1.0, - cost_fn: Optional[costs.CostFn] = None, + cost_fn: costs.CostFn | None = None, **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Compute solution to a quadratic OT problem. Args: @@ -95,7 +95,7 @@ def match_quadratic( def sample_joint(rng: jax.Array, - tmat: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]: + tmat: jax.Array) -> tuple[jax.Array, jax.Array]: """Sample jointly from a transport matrix. Args: @@ -117,10 +117,10 @@ def sample_joint(rng: jax.Array, def sample_conditional( rng: jax.Array, - tmat: jnp.ndarray, + tmat: jax.Array, *, k: int = 1, -) -> Tuple[jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array]: """Sample conditionally from a transport matrix. Args: @@ -154,8 +154,8 @@ def uniform_sampler( num_samples: int, low: float = 0.0, high: float = 1.0, - offset: Optional[float] = None -) -> jnp.ndarray: + offset: float | None = None +) -> jax.Array: r"""Sample from a uniform distribution. Sample :math:`t` from a uniform distribution :math:`[low, high]`. diff --git a/src/ott/solvers/was_solver.py b/src/ott/solvers/was_solver.py index d5ed178a1..fd61200d9 100644 --- a/src/ott/solvers/was_solver.py +++ b/src/ott/solvers/was_solver.py @@ -11,7 +11,8 @@ # 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 typing import TYPE_CHECKING, Any, Dict, Sequence, Tuple, Union +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any, Union import jax import jax.numpy as jnp @@ -33,7 +34,7 @@ class WassersteinSolver: def __init__( self, - linear_solver: Union["sinkhorn.Sinkhorn", "sinkhorn_lr.LRSinkhorn"], + linear_solver: sinkhorn.Sinkhorn | sinkhorn_lr.LRSinkhorn, threshold: float = 1e-3, min_iterations: int = 5, max_iterations: int = 50, @@ -55,7 +56,7 @@ def is_low_rank(self) -> bool: """Whether the solver is low-rank.""" return isinstance(self.linear_solver, sinkhorn_lr.LRSinkhorn) - def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 + def tree_flatten(self) -> tuple[Sequence[Any], dict[str, Any]]: # noqa: D102 return ([self.linear_solver, self.threshold], { "min_iterations": self.min_iterations, "max_iterations": self.max_iterations, @@ -64,7 +65,7 @@ def tree_flatten(self) -> Tuple[Sequence[Any], Dict[str, Any]]: # noqa: D102 @classmethod def tree_unflatten( # noqa: D102 - cls, aux_data: Dict[str, Any], children: Sequence[Any] + cls, aux_data: dict[str, Any], children: Sequence[Any] ) -> "WassersteinSolver": return cls(*children, **aux_data) diff --git a/src/ott/tools/conformal.py b/src/ott/tools/conformal.py index c75475ed1..d44c05c5a 100644 --- a/src/ott/tools/conformal.py +++ b/src/ott/tools/conformal.py @@ -14,13 +14,15 @@ import dataclasses import math import operator -from typing import Any, Callable, Optional, Tuple +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp import jax.tree_util as jtu import numpy as np import scipy as sp +from jax.typing import ArrayLike from scipy.stats import qmc from ott import utils @@ -30,14 +32,14 @@ __all__ = ["OTCP", "sobol_ball_sampler"] -ScoreFn = Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray] +ScoreFn = Callable[[jax.Array, jax.Array], jax.Array] def sobol_ball_sampler( - rng: Optional[jax.Array], - shape: Tuple[int, int], - n_per_radius: Optional[int] = None, -) -> Tuple[jnp.ndarray, jnp.ndarray]: + rng: jax.Array | None, + shape: tuple[int, int], + n_per_radius: int | None = None, +) -> tuple[jax.Array, jax.Array]: """Sample target measure for :class:`OTCP`. Args: @@ -103,28 +105,28 @@ class OTCP: calibration_scores: Nonconformity calibration scores computed in :meth:`calibrate`. """ - model: Callable[[jnp.ndarray], - jnp.ndarray] = dataclasses.field(metadata={"static": True}) + model: Callable[[jax.Array], + jax.Array] = dataclasses.field(metadata={"static": True}) nonconformity_fn: ScoreFn = dataclasses.field( default=operator.sub, metadata={"static": True} ) - sinkhorn_output: Optional[sinkhorn.SinkhornOutput] = None - sampler: Optional[Callable[[jax.random.PRNGKey, Tuple[int, int]], - jax.Array]] = dataclasses.field( - default=None, metadata={"static": True} - ) - offset: jnp.ndarray = 0.0 - scale: jnp.ndarray = 1.0 - calibration_scores: Optional[jnp.ndarray] = None + sinkhorn_output: sinkhorn.SinkhornOutput | None = None + sampler: Callable[[jax.Array, tuple[int, int]], + jax.Array] | None = dataclasses.field( + default=None, metadata={"static": True} + ) + offset: ArrayLike = 0.0 + scale: ArrayLike = 1.0 + calibration_scores: jax.Array | None = None def fit_transport( self, - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, epsilon: float = 1e-1, n_target: int = 8192, - rng: Optional[jax.Array] = None, - sampler_kwargs: Optional[Any] = None, + rng: jax.Array | None = None, + sampler_kwargs: Any | None = None, **kwargs: Any, ) -> "OTCP": """Fit the transport map. @@ -163,7 +165,7 @@ def fit_transport( self, sinkhorn_output=out, offset=offset, scale=scale ) - def calibrate(self, x: jnp.ndarray, y: jnp.ndarray) -> "OTCP": + def calibrate(self, x: jax.Array, y: jax.Array) -> "OTCP": """Compute calibration scores. Args: @@ -176,7 +178,7 @@ def calibrate(self, x: jnp.ndarray, y: jnp.ndarray) -> "OTCP": scores = self.get_scores(x, y) return dataclasses.replace(self, calibration_scores=scores) - def get_scores(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + def get_scores(self, x: jax.Array, y: jax.Array) -> jax.Array: """Compute nonconformity scores. Args: @@ -190,10 +192,10 @@ def get_scores(self, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: def predict( self, - x: jnp.ndarray, - y_candidates: Optional[jnp.ndarray] = None, + x: jax.Array, + y_candidates: jax.Array | None = None, alpha: float = 0.1, - ) -> jnp.ndarray: + ) -> jax.Array: """Conformalize the model's prediction. Args: @@ -218,21 +220,21 @@ def predict( def _predict_backward( self, - y_hat: jnp.ndarray, + y_hat: jax.Array, *, quantile: float, - ) -> jnp.ndarray: + ) -> jax.Array: candidates = self._transport(quantile * self.target_measure, forward=False) candidates = self._rescale(candidates, forward=False) return y_hat[:, None] + candidates[None] def _predict_forward( self, - y_hat: jnp.ndarray, - y_candidates: jnp.ndarray, + y_hat: jax.Array, + y_candidates: jax.Array, *, quantile: float, - ) -> jnp.ndarray: + ) -> jax.Array: assert y_candidates.ndim == 2, y_candidates.shape score_fn = jax.vmap( jax.vmap(self._get_scores, in_axes=[0, None]), in_axes=[None, 0] @@ -240,25 +242,25 @@ def _predict_forward( scores = score_fn(y_candidates, y_hat) return scores <= quantile - def _get_scores(self, y: jnp.ndarray, y_hat: jnp.ndarray) -> jnp.ndarray: + def _get_scores(self, y: jax.Array, y_hat: jax.Array) -> jax.Array: scores = self.nonconformity_fn(jnp.atleast_2d(y), jnp.atleast_2d(y_hat)) scores = self._rescale(scores, forward=True) scores = self._transport(scores, forward=True) scores = jnp.linalg.norm(scores, axis=-1) return scores.squeeze(0) if y.ndim == 1 else scores - def _transport(self, x: jnp.ndarray, *, forward: bool = True) -> jnp.ndarray: + def _transport(self, x: jax.Array, *, forward: bool = True) -> jax.Array: assert self.sinkhorn_output is not None, "Run `.fit_transport()` first." return self.sinkhorn_output.to_dual_potentials().transport( x, forward=forward ) - def _rescale(self, x: jnp.ndarray, *, forward: bool) -> jnp.ndarray: + def _rescale(self, x: jax.Array, *, forward: bool) -> jax.Array: if forward: return (x - self.offset) / self.scale return (self.scale * x) + self.offset @property - def target_measure(self) -> Optional[jnp.ndarray]: + def target_measure(self) -> jax.Array | None: """Target measure of shape ``[n_target, dim_y]``.""" return None if self.sinkhorn_output is None else self.sinkhorn_output.geom.y diff --git a/src/ott/tools/gaussian_mixture/fit_gmm.py b/src/ott/tools/gaussian_mixture/fit_gmm.py index 6b5cdfe54..0b66f7f30 100644 --- a/src/ott/tools/gaussian_mixture/fit_gmm.py +++ b/src/ott/tools/gaussian_mixture/fit_gmm.py @@ -49,8 +49,6 @@ $$ """ -from typing import Optional - import jax import jax.numpy as jnp @@ -62,8 +60,8 @@ def get_assignment_probs( - gmm: gaussian_mixture.GaussianMixture, points: jnp.ndarray -) -> jnp.ndarray: + gmm: gaussian_mixture.GaussianMixture, points: jax.Array +) -> jax.Array: r"""Get component assignment probabilities used in the E step of EM. Here we compute the component assignment probabilities p(Z|X, \Theta^{(t)}) @@ -81,9 +79,9 @@ def get_assignment_probs( def get_q( gmm: gaussian_mixture.GaussianMixture, - assignment_probs: jnp.ndarray, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray] = None, + assignment_probs: jax.Array, + points: jax.Array, + point_weights: jax.Array | None = None, ) -> float: r"""Get Q(\Theta|\Theta^{(t)}). @@ -109,8 +107,8 @@ def get_q( def log_prob_loss( gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray] = None, + points: jax.Array, + point_weights: jax.Array | None = None, ) -> float: """Loss function: weighted mean of (-log prob of observations). @@ -130,8 +128,8 @@ def log_prob_loss( def fit_model_em( gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], + points: jax.Array, + point_weights: jax.Array | None, steps: int, jit: bool = True, verbose: bool = False, @@ -184,10 +182,10 @@ def fit_model_em( # See https://en.wikipedia.org/wiki/K-means%2B%2B for details -def _get_dist_sq(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: +def _get_dist_sq(points: jax.Array, loc: jax.Array) -> jax.Array: """Get the squared distance from each point to each loc.""" - def _dist_sq_one_loc(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: + def _dist_sq_one_loc(points: jax.Array, loc: jax.Array) -> jax.Array: return jnp.sum((points - loc[None]) ** 2, axis=-1) dist_sq_fn = jax.vmap(_dist_sq_one_loc, in_axes=(None, 0), out_axes=1) @@ -195,8 +193,8 @@ def _dist_sq_one_loc(points: jnp.ndarray, loc: jnp.ndarray) -> jnp.ndarray: def _get_locs( - rng: jax.Array, points: jnp.ndarray, n_components: int -) -> jnp.ndarray: + rng: jax.Array, points: jax.Array, n_components: int +) -> jax.Array: """Get the initial component means. Args: @@ -230,8 +228,8 @@ def _get_locs( def from_kmeans_plusplus( rng: jax.Array, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], + points: jax.Array, + point_weights: jax.Array | None, n_components: int, ) -> gaussian_mixture.GaussianMixture: """Initialize a GMM via a single pass of K-means++. @@ -266,8 +264,8 @@ def from_kmeans_plusplus( def initialize( rng: jax.Array, - points: jnp.ndarray, - point_weights: Optional[jnp.ndarray], + points: jax.Array, + point_weights: jax.Array | None, n_components: int, n_attempts: int = 50, verbose: bool = False diff --git a/src/ott/tools/gaussian_mixture/fit_gmm_pair.py b/src/ott/tools/gaussian_mixture/fit_gmm_pair.py index 084e9166f..4934c3cbe 100644 --- a/src/ott/tools/gaussian_mixture/fit_gmm_pair.py +++ b/src/ott/tools/gaussian_mixture/fit_gmm_pair.py @@ -79,7 +79,8 @@ import functools import math -from typing import Callable, NamedTuple, Optional, Tuple +from collections.abc import Callable +from typing import NamedTuple import jax import jax.numpy as jnp @@ -98,9 +99,9 @@ class Observations(NamedTuple): """Weighted observations and their E-step assignment probabilities.""" - points: jnp.ndarray - point_weights: jnp.ndarray - assignment_probs: jnp.ndarray + points: jax.Array + point_weights: jax.Array + assignment_probs: jax.Array # Model fit @@ -108,7 +109,7 @@ class Observations(NamedTuple): def get_q( gmm: gaussian_mixture.GaussianMixture, obs: Observations -) -> jnp.ndarray: +) -> jax.Array: r"""Get Q(\Theta|\Theta^{(t)}). Here Q is the log likelihood for our observations based on the current @@ -159,7 +160,7 @@ def _objective_fn( pair: gaussian_mixture_pair.GaussianMixturePair, obs0: Observations, obs1: Observations, - ) -> jnp.ndarray: + ) -> jax.Array: """Compute the objective function for a pair of GMMs. Args: @@ -204,11 +205,11 @@ def print_losses( def do_e_step( # noqa: D103 - e_step_fn: Callable[[gaussian_mixture.GaussianMixture, jnp.ndarray], - jnp.ndarray], + e_step_fn: Callable[[gaussian_mixture.GaussianMixture, jax.Array], + jax.Array], gmm: gaussian_mixture.GaussianMixture, - points: jnp.ndarray, - point_weights: jnp.ndarray, + points: jax.Array, + point_weights: jax.Array, ) -> Observations: assignment_probs = e_step_fn(gmm, points) return Observations( @@ -307,14 +308,14 @@ def get_fit_model_em_fn( def _fit_model_em( pair: gaussian_mixture_pair.GaussianMixturePair, - points0: jnp.ndarray, - points1: jnp.ndarray, - point_weights0: Optional[jnp.ndarray], - point_weights1: Optional[jnp.ndarray], + points0: jax.Array, + points1: jax.Array, + point_weights0: jax.Array | None, + point_weights1: jax.Array | None, em_steps: int, m_steps: int = 50, verbose: bool = False, - ) -> Tuple[gaussian_mixture_pair.GaussianMixturePair, float]: + ) -> tuple[gaussian_mixture_pair.GaussianMixturePair, float]: """Optimize a GaussianMixturePair using penalized EM. Args: diff --git a/src/ott/tools/gaussian_mixture/gaussian.py b/src/ott/tools/gaussian_mixture/gaussian.py index d54f92e8b..032a8b571 100644 --- a/src/ott/tools/gaussian_mixture/gaussian.py +++ b/src/ott/tools/gaussian_mixture/gaussian.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. import math -from typing import Optional, Union import jax import jax.numpy as jnp +from jax.typing import ArrayLike from ott.tools.gaussian_mixture import scale_tril @@ -28,15 +28,13 @@ class Gaussian: """Normal distribution.""" - def __init__(self, loc: jnp.ndarray, scale: scale_tril.ScaleTriL): + def __init__(self, loc: jax.Array, scale: scale_tril.ScaleTriL): self._loc = loc self._scale = scale @classmethod def from_samples( - cls, - points: jnp.ndarray, - weights: Optional[jnp.ndarray] = None + cls, points: jax.Array, weights: jax.Array | None = None ) -> "Gaussian": """Construct a Gaussian from weighted samples. @@ -67,7 +65,7 @@ def from_random( n_dimensions: int, stdev_mean: float = 0.1, stdev_cov: float = 0.1, - ridge: Union[float, jnp.ndarray] = 0, + ridge: ArrayLike = 0, ) -> "Gaussian": """Construct a random Gaussian. @@ -90,13 +88,13 @@ def from_random( return cls(loc=loc, scale=scale) @classmethod - def from_mean_and_cov(cls, mean: jnp.ndarray, cov: jnp.ndarray) -> "Gaussian": + def from_mean_and_cov(cls, mean: jax.Array, cov: jax.Array) -> "Gaussian": """Construct a Gaussian from a mean and covariance.""" scale = scale_tril.ScaleTriL.from_covariance(cov) return cls(loc=mean, scale=scale) @property - def loc(self) -> jnp.ndarray: + def loc(self) -> jax.Array: """Mean of the Gaussian.""" return self._loc @@ -110,22 +108,22 @@ def n_dimensions(self) -> int: """Dimensionality of the Gaussian.""" return self.loc.shape[-1] - def covariance(self) -> jnp.ndarray: + def covariance(self) -> jax.Array: """Covariance of the Gaussian.""" return self.scale.covariance() - def to_z(self, x: jnp.ndarray) -> jnp.ndarray: + def to_z(self, x: jax.Array) -> jax.Array: r"""Transform :math:`x` to :math:`z = \frac{x - loc}{scale}`.""" return self.scale.centered_to_z(x_centered=x - self.loc) - def from_z(self, z: jnp.ndarray) -> jnp.ndarray: + def from_z(self, z: jax.Array) -> jax.Array: r"""Transform :math:`z` to :math:`x = loc + scale \cdot z`.""" return self.scale.z_to_centered(z=z) + self.loc def log_prob( self, - x: jnp.ndarray, # (?, d) - ) -> jnp.ndarray: # (?, d) + x: jax.Array, # (?, d) + ) -> jax.Array: # (?, d) """Log probability for a Gaussian with a diagonal covariance.""" d = x.shape[-1] z = self.to_z(x) @@ -134,7 +132,7 @@ def log_prob( -0.5 * (d * LOG2PI + log_det[None] + jnp.sum(z ** 2, axis=-1)) ) # (?, k) - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + def sample(self, rng: jax.Array, size: int) -> jax.Array: """Generate samples from the distribution.""" std_samples_t = jax.random.normal(rng, shape=(self.n_dimensions, size)) return self.loc[None] + ( @@ -145,7 +143,7 @@ def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: ) ) - def w2_dist(self, other: "Gaussian") -> jnp.ndarray: + def w2_dist(self, other: "Gaussian") -> jax.Array: r"""Wasserstein distance :math:`W_2^2` to another Gaussian. .. math:: @@ -163,7 +161,7 @@ def w2_dist(self, other: "Gaussian") -> jnp.ndarray: delta_sigma = self.scale.w2_dist(other.scale) return delta_mean + delta_sigma - def f_potential(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: + def f_potential(self, dest: "Gaussian", points: jax.Array) -> jax.Array: """Optimal potential for W2 distance between Gaussians. Evaluated on points. Args: @@ -187,7 +185,7 @@ def batch_inner_product(x, y): points.dot(dest.loc) ) - def transport(self, dest: "Gaussian", points: jnp.ndarray) -> jnp.ndarray: + def transport(self, dest: "Gaussian", points: jax.Array) -> jax.Array: """Transport points according to map between two Gaussian measures. Args: @@ -209,9 +207,3 @@ def tree_flatten(self): # noqa: D102 @classmethod def tree_unflatten(cls, aux_data, children): # noqa: D102 return cls(*children, **aux_data) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/src/ott/tools/gaussian_mixture/gaussian_mixture.py b/src/ott/tools/gaussian_mixture/gaussian_mixture.py index 42925acf3..3cde75d88 100644 --- a/src/ott/tools/gaussian_mixture/gaussian_mixture.py +++ b/src/ott/tools/gaussian_mixture/gaussian_mixture.py @@ -11,7 +11,6 @@ # 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 typing import List, Tuple, Union import jax import jax.numpy as jnp @@ -27,9 +26,8 @@ def get_summary_stats_from_points_and_assignment_probs( - points: jnp.ndarray, point_weights: jnp.ndarray, - assignment_probs: jnp.ndarray -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: + points: jax.Array, point_weights: jax.Array, assignment_probs: jax.Array +) -> tuple[jax.Array, jax.Array, jax.Array]: """Get component summary stats from points and component probabilities. Args: @@ -68,7 +66,7 @@ class GaussianMixture: """Gaussian Mixture model.""" def __init__( - self, loc: jnp.ndarray, scale_params: jnp.ndarray, + self, loc: jax.Array, scale_params: jax.Array, component_weight_ob: probabilities.Probabilities ): self._loc = loc @@ -84,7 +82,7 @@ def from_random( stdev_mean: float = 0.1, stdev_cov: float = 0.1, stdev_weights: float = 0.1, - ridge: Union[float, jnp.array] = 0, + ridge: float | jnp.array = 0, ) -> "GaussianMixture": """Construct a random GMM.""" loc = [] @@ -113,7 +111,7 @@ def from_random( @classmethod def from_mean_cov_component_weights( - cls, mean: jnp.ndarray, cov: jnp.ndarray, component_weights: jnp.ndarray + cls, mean: jax.Array, cov: jax.Array, component_weights: jax.Array ): """Construct a GMM from means, covariances, and component weights.""" scale_params = [] @@ -128,9 +126,9 @@ def from_mean_cov_component_weights( @classmethod def from_points_and_assignment_probs( cls, - points: jnp.ndarray, - point_weights: jnp.ndarray, - assignment_probs: jnp.ndarray, + points: jax.Array, + point_weights: jax.Array, + assignment_probs: jax.Array, ) -> "GaussianMixture": """Estimate a GMM from points and a set of component probabilities.""" mean, cov, wts = get_summary_stats_from_points_and_assignment_probs( @@ -158,17 +156,17 @@ def n_components(self): return self._loc.shape[-2] @property - def loc(self) -> jnp.ndarray: + def loc(self) -> jax.Array: """Location parameters of the GMM.""" return self._loc @property - def scale_params(self) -> jnp.ndarray: + def scale_params(self) -> jax.Array: """Scale parameters of the GMM.""" return self._scale_params @property - def cholesky(self) -> jnp.ndarray: + def cholesky(self) -> jax.Array: """Cholesky decomposition of the GMM covariance matrices.""" size = self.n_dimensions @@ -178,7 +176,7 @@ def _get_cholesky(scale_params): return jax.vmap(_get_cholesky, in_axes=0, out_axes=0)(self.scale_params) @property - def covariance(self) -> jnp.ndarray: + def covariance(self) -> jax.Array: """Covariance matrices of the GMM.""" size = self.n_dimensions @@ -193,16 +191,16 @@ def component_weight_ob(self) -> probabilities.Probabilities: return self._component_weight_ob @property - def component_weights(self) -> jnp.ndarray: + def component_weights(self) -> jax.Array: """Component weights probabilities.""" return self._component_weight_ob.probs() - def log_component_weights(self) -> jnp.ndarray: + def log_component_weights(self) -> jax.Array: """Log component weights probabilities.""" return self._component_weight_ob.log_probs() def _get_normal( - self, loc: jnp.ndarray, scale_params: jnp.ndarray + self, loc: jax.Array, scale_params: jax.Array ) -> gaussian.Gaussian: size = loc.shape[-1] return gaussian.Gaussian( @@ -215,11 +213,11 @@ def get_component(self, index: int) -> gaussian.Gaussian: loc=self.loc[index], scale_params=self.scale_params[index] ) - def components(self) -> List[gaussian.Gaussian]: + def components(self) -> list[gaussian.Gaussian]: """List of all GMM components.""" return [self.get_component(i) for i in range(self.n_components)] - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + def sample(self, rng: jax.Array, size: int) -> jax.Array: """Generate samples from the distribution.""" subrng0, subrng1 = jax.random.split(rng) component = self.component_weight_ob.sample(rng=subrng0, size=size) @@ -242,7 +240,7 @@ def _transform_single_value(single_component, single_x): axis=0 ) - def conditional_log_prob(self, x: jnp.ndarray) -> jnp.ndarray: + def conditional_log_prob(self, x: jax.Array) -> jax.Array: """Compute the component-conditional log probability of x. Args: @@ -254,7 +252,7 @@ def conditional_log_prob(self, x: jnp.ndarray) -> jnp.ndarray: """ def _log_prob_single_component( - loc: jnp.ndarray, scale_params: jnp.ndarray, x: jnp.ndarray + loc: jax.Array, scale_params: jax.Array, x: jax.Array ): norm = self._get_normal(loc=loc, scale_params=scale_params) return norm.log_prob(x) @@ -264,7 +262,7 @@ def _log_prob_single_component( ) return conditional_log_prob_fn(self._loc, self._scale_params, x) - def log_prob(self, x: jnp.ndarray) -> jnp.ndarray: + def log_prob(self, x: jax.Array) -> jax.Array: """Compute the log probability of the observations x. Args: @@ -280,7 +278,7 @@ def log_prob(self, x: jnp.ndarray) -> jnp.ndarray: log_prob_conditional + log_component_weight[None, :], axis=-1 ) - def get_log_component_posterior(self, x: jnp.ndarray) -> jnp.ndarray: + def get_log_component_posterior(self, x: jax.Array) -> jax.Array: """Compute the posterior probability that x came from each component. Args: @@ -300,10 +298,7 @@ def get_log_component_posterior(self, x: jnp.ndarray) -> jnp.ndarray: ) def has_nans(self) -> bool: # noqa: D102 - for leaf in jax.tree_util.tree_leaves(self): - if jnp.any(~jnp.isfinite(leaf)): - return True - return False + return any(jnp.any(~jnp.isfinite(leaf)) for leaf in jax.tree.leaves(self)) def tree_flatten(self): # noqa: D102 children = (self.loc, self.scale_params, self.component_weight_ob) @@ -321,9 +316,3 @@ def __repr__(self): class_name, ", ".join([repr(c) for c in children] + [f"{k}: {repr(v)}" for k, v in aux.items()]) ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py b/src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py index 2b250f80c..809d3c903 100644 --- a/src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py +++ b/src/ott/tools/gaussian_mixture/gaussian_mixture_pair.py @@ -128,12 +128,12 @@ def get_bures_geometry(self) -> pointcloud.PointCloud: epsilon=self.epsilon ) - def get_cost_matrix(self) -> jnp.ndarray: + def get_cost_matrix(self) -> jax.Array: """Get matrix of :math:`W_2^2` costs between all pairs of components.""" return self.get_bures_geometry().cost_matrix def get_sinkhorn( - self, cost_matrix: jnp.ndarray, **kwargs: Any + self, cost_matrix: jax.Array, **kwargs: Any ) -> sinkhorn.SinkhornOutput: """Get the output of Sinkhorn's method for a given cost matrix.""" # We use a Geometry here rather than the PointCloud created in @@ -152,7 +152,7 @@ def get_sinkhorn( def get_normalized_sinkhorn_coupling( self, sinkhorn_output: sinkhorn.SinkhornOutput, - ) -> jnp.ndarray: + ) -> jax.Array: """Get the normalized coupling matrix for the specified Sinkhorn output. Args: @@ -214,9 +214,3 @@ def __repr__(self): class_name, ", ".join([repr(c) for c in children] + [f"{k}: {repr(v)}" for k, v in aux.items()]) ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/src/ott/tools/gaussian_mixture/linalg.py b/src/ott/tools/gaussian_mixture/linalg.py index 885b2e508..6e314ed98 100644 --- a/src/ott/tools/gaussian_mixture/linalg.py +++ b/src/ott/tools/gaussian_mixture/linalg.py @@ -11,16 +11,16 @@ # 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 typing import Callable, Iterable, List, Tuple +from collections.abc import Callable, Iterable import jax import jax.numpy as jnp def get_mean_and_var( - points: jnp.ndarray, # (n, d) - weights: jnp.ndarray, # (n,) -) -> Tuple[jnp.ndarray, jnp.ndarray]: + points: jax.Array, # (n, d) + weights: jax.Array, # (n,) +) -> tuple[jax.Array, jax.Array]: """Get the mean and variance of a weighted set of points.""" weights_sum = jnp.sum(weights, axis=-1) # (1,) mean = ( @@ -37,9 +37,9 @@ def get_mean_and_var( def get_mean_and_cov( - points: jnp.ndarray, # (n, d) - weights: jnp.ndarray, # (n,) -) -> Tuple[jnp.ndarray, jnp.ndarray]: + points: jax.Array, # (n, d) + weights: jax.Array, # (n,) +) -> tuple[jax.Array, jax.Array]: """Get the mean and covariance of a weighted set of points.""" weights_sum = jnp.sum(weights, axis=-1, keepdims=True) # (1,) mean = ( @@ -59,7 +59,7 @@ def get_mean_and_cov( return mean, cov -def flat_to_tril(x: jnp.ndarray, size: int) -> jnp.ndarray: +def flat_to_tril(x: jax.Array, size: int) -> jax.Array: """Map flat values to lower triangular matrices. Args: @@ -76,7 +76,7 @@ def flat_to_tril(x: jnp.ndarray, size: int) -> jnp.ndarray: return m.at[..., tril[0], tril[1]].set(x) -def tril_to_flat(m: jnp.ndarray) -> jnp.ndarray: +def tril_to_flat(m: jax.Array) -> jax.Array: """Flatten lower triangular matrices. Args: @@ -91,8 +91,8 @@ def tril_to_flat(m: jnp.ndarray) -> jnp.ndarray: def apply_to_diag( - m: jnp.ndarray, fn: Callable[[jnp.ndarray], jnp.ndarray] -) -> jnp.ndarray: + m: jax.Array, fn: Callable[[jax.Array], jax.Array] +) -> jax.Array: """Apply a function to the diagonal of a matrix.""" size = m.shape[-1] diag = jnp.diagonal(m, axis1=-2, axis2=-1) @@ -101,9 +101,9 @@ def apply_to_diag( def matrix_powers( - m: jnp.ndarray, + m: jax.Array, powers: Iterable[float], -) -> List[jnp.ndarray]: +) -> list[jax.Array]: """Raise a real, symmetric matrix to multiple powers.""" eigs, q = jnp.linalg.eigh(m) qt = jnp.swapaxes(q, axis1=-2, axis2=-1) @@ -113,9 +113,7 @@ def matrix_powers( return ret -def invmatvectril( - m: jnp.ndarray, x: jnp.ndarray, lower: bool = True -) -> jnp.ndarray: +def invmatvectril(m: jax.Array, x: jax.Array, lower: bool = True) -> jax.Array: """Multiply x by the inverse of a triangular matrix. Args: @@ -131,7 +129,7 @@ def invmatvectril( ) -def get_random_orthogonal(rng: jax.Array, dim: int) -> jnp.ndarray: +def get_random_orthogonal(rng: jax.Array, dim: int) -> jax.Array: """Get a random orthogonal matrix with the specified dimension.""" m = jax.random.normal(rng, shape=[dim, dim]) q, _ = jnp.linalg.qr(m) diff --git a/src/ott/tools/gaussian_mixture/probabilities.py b/src/ott/tools/gaussian_mixture/probabilities.py index 2587235c9..147d6187f 100644 --- a/src/ott/tools/gaussian_mixture/probabilities.py +++ b/src/ott/tools/gaussian_mixture/probabilities.py @@ -11,7 +11,6 @@ # 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 typing import Optional import jax import jax.numpy as jnp @@ -27,7 +26,7 @@ class Probabilities: to a length n simplex by appending a 0 and taking a softmax. """ - _params: jnp.ndarray + _params: jax.Array def __init__(self, params): self._params = params @@ -37,13 +36,13 @@ def from_random( cls, rng: jax.Array, n_dimensions: int, - stdev: Optional[float] = 0.1, + stdev: float | None = 0.1, ) -> "Probabilities": """Construct a random Probabilities.""" return cls(params=jax.random.normal(rng, shape=(n_dimensions - 1,)) * stdev) @classmethod - def from_probs(cls, probs: jnp.ndarray) -> "Probabilities": + def from_probs(cls, probs: jax.Array) -> "Probabilities": """Construct Probabilities from a vector of probabilities.""" log_probs = jnp.log(probs) log_probs_normalized, norm = log_probs[:-1], log_probs[-1] @@ -58,19 +57,19 @@ def params(self): # noqa: D102 def dtype(self): # noqa: D102 return self._params.dtype - def unnormalized_log_probs(self) -> jnp.ndarray: + def unnormalized_log_probs(self) -> jax.Array: """Get the unnormalized log probabilities.""" return jnp.concatenate([self._params, jnp.zeros((1,))], axis=-1) - def log_probs(self) -> jnp.ndarray: + def log_probs(self) -> jax.Array: """Get the log probabilities.""" return jax.nn.log_softmax(self.unnormalized_log_probs()) - def probs(self) -> jnp.ndarray: + def probs(self) -> jax.Array: """Get the probabilities.""" return jax.nn.softmax(self.unnormalized_log_probs()) - def sample(self, rng: jax.Array, size: int) -> jnp.ndarray: + def sample(self, rng: jax.Array, size: int) -> jax.Array: """Sample from the distribution.""" return jax.random.categorical( rng, logits=self.unnormalized_log_probs(), shape=(size,) @@ -92,9 +91,3 @@ def __repr__(self): class_name, ", ".join([repr(c) for c in children] + [f"{k}: {repr(v)}" for k, v in aux.items()]) ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/src/ott/tools/gaussian_mixture/scale_tril.py b/src/ott/tools/gaussian_mixture/scale_tril.py index 5da706690..596d9008f 100644 --- a/src/ott/tools/gaussian_mixture/scale_tril.py +++ b/src/ott/tools/gaussian_mixture/scale_tril.py @@ -11,7 +11,6 @@ # 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 typing import Optional, Tuple import jax import jax.numpy as jnp @@ -27,16 +26,16 @@ class ScaleTriL: """Pytree for a lower triangular Cholesky-factored covariance matrix.""" - def __init__(self, params: jnp.ndarray, size: int): + def __init__(self, params: jax.Array, size: int): self._params = params self._size = size @classmethod def from_points_and_weights( cls, - points: jnp.ndarray, - weights: jnp.ndarray, - ) -> Tuple[jnp.ndarray, "ScaleTriL"]: + points: jax.Array, + weights: jax.Array, + ) -> tuple[jax.Array, "ScaleTriL"]: """Get a mean and a ScaleTriL from a set of points and weights.""" mean, cov = linalg.get_mean_and_cov(points=points, weights=weights) return mean, cls.from_covariance(cov) @@ -46,7 +45,7 @@ def from_random( cls, rng: jax.Array, n_dimensions: int, - stdev: Optional[float] = 0.1, + stdev: float | None = 0.1, ) -> "ScaleTriL": """Construct a random ScaleTriL. @@ -76,7 +75,7 @@ def from_random( return cls(params=flat, size=n_dimensions) @classmethod - def from_cholesky(cls, cholesky: jnp.ndarray) -> "ScaleTriL": + def from_cholesky(cls, cholesky: jax.Array) -> "ScaleTriL": """Construct ScaleTriL from a Cholesky factor of a covariance matrix.""" m = linalg.apply_to_diag(cholesky, jnp.log) flat = linalg.tril_to_flat(m) @@ -85,14 +84,14 @@ def from_cholesky(cls, cholesky: jnp.ndarray) -> "ScaleTriL": @classmethod def from_covariance( cls, - covariance: jnp.ndarray, + covariance: jax.Array, ) -> "ScaleTriL": """Construct ScaleTriL from a covariance matrix.""" cholesky = jnp.linalg.cholesky(covariance) return cls.from_cholesky(cholesky) @property - def params(self) -> jnp.ndarray: + def params(self) -> jax.Array: """Internal representation.""" return self._params @@ -106,34 +105,34 @@ def dtype(self): """Data type of the covariance matrix.""" return self._params.dtype - def cholesky(self) -> jnp.ndarray: + def cholesky(self) -> jax.Array: """Get a lower triangular Cholesky factor for the covariance matrix.""" m = linalg.flat_to_tril(self._params, size=self._size) return linalg.apply_to_diag(m, jnp.exp) - def covariance(self) -> jnp.ndarray: + def covariance(self) -> jax.Array: """Get the covariance matrix.""" cholesky = self.cholesky() return cholesky @ cholesky.T - def covariance_sqrt(self) -> jnp.ndarray: + def covariance_sqrt(self) -> jax.Array: """Get the square root of the covariance matrix.""" return linalg.matrix_powers(self.covariance(), (0.5,))[0] - def log_det_covariance(self) -> jnp.ndarray: + def log_det_covariance(self) -> jax.Array: """Get the log of the determinant of the covariance matrix.""" diag = jnp.diagonal(self.cholesky(), axis1=-2, axis2=-1) return 2.0 * jnp.sum(jnp.log(diag), axis=-1) - def centered_to_z(self, x_centered: jnp.ndarray) -> jnp.ndarray: + def centered_to_z(self, x_centered: jax.Array) -> jax.Array: """Map centered points to standardized centered points (i.e. cov(z) = I).""" return linalg.invmatvectril(m=self.cholesky(), x=x_centered, lower=True) - def z_to_centered(self, z: jnp.ndarray) -> jnp.ndarray: + def z_to_centered(self, z: jax.Array) -> jax.Array: """Scale standardized points to points with the specified covariance.""" return (self.cholesky() @ z.T).T - def w2_dist(self, other: "ScaleTriL") -> jnp.ndarray: + def w2_dist(self, other: "ScaleTriL") -> jax.Array: r"""Wasserstein distance W_2^2 to another Gaussian with same mean. Args: @@ -144,7 +143,7 @@ def w2_dist(self, other: "ScaleTriL") -> jnp.ndarray: """ dimension = self.size - def _flatten_cov(cov: jnp.ndarray) -> jnp.ndarray: + def _flatten_cov(cov: jax.Array) -> jax.Array: cov = cov.reshape(cov.shape[:-2] + (dimension * dimension,)) return jnp.concatenate([jnp.zeros(dimension), cov], axis=-1) @@ -153,7 +152,7 @@ def _flatten_cov(cov: jnp.ndarray) -> jnp.ndarray: cost_fn = costs.Bures(dimension=dimension) return cost_fn(x0, x1) - def gaussian_map(self, dest_scale: "ScaleTriL") -> jnp.ndarray: + def gaussian_map(self, dest_scale: "ScaleTriL") -> jax.Array: """Scaling matrix used in transport between 0-mean Gaussians. Sigma_mu^{-1/2} @ @@ -173,9 +172,7 @@ def gaussian_map(self, dest_scale: "ScaleTriL") -> jnp.ndarray: ) return jnp.matmul(sqrt0_inv, jnp.matmul(m, sqrt0_inv)) - def transport( - self, dest_scale: "ScaleTriL", points: jnp.ndarray - ) -> jnp.ndarray: + def transport(self, dest_scale: "ScaleTriL", points: jax.Array) -> jax.Array: """Apply Monge map, computed between two 0-mean Gaussians, to points. Args: @@ -204,9 +201,3 @@ def __repr__(self): class_name, ", ".join([repr(c) for c in children] + [f"{k}: {repr(v)}" for k, v in aux.items()]) ) - - def __hash__(self): - return jax.tree_util.tree_flatten(self).__hash__() - - def __eq__(self, other): - return jax.tree_util.tree_flatten(self) == jax.tree_util.tree_flatten(other) diff --git a/src/ott/tools/k_means.py b/src/ott/tools/k_means.py index 2b54679a0..d65269b04 100644 --- a/src/ott/tools/k_means.py +++ b/src/ott/tools/k_means.py @@ -13,7 +13,8 @@ # limitations under the License. import functools import math -from typing import Callable, Literal, NamedTuple, Optional, Tuple, Union +from collections.abc import Callable +from typing import Literal, NamedTuple import jax import jax.numpy as jnp @@ -24,30 +25,31 @@ __all__ = ["k_means", "KMeansOutput"] -Init_t = Union[Literal["k-means++", "random"], - Callable[[pointcloud.PointCloud, int, jnp.ndarray], jnp.ndarray]] +Init_t = Literal["k-means++", + "random"] | Callable[[pointcloud.PointCloud, int, jax.Array], + jax.Array] class KPPState(NamedTuple): # noqa: D101 rng: jax.Array - centroids: jnp.ndarray - centroid_dists: jnp.ndarray + centroids: jax.Array + centroid_dists: jax.Array class KMeansState(NamedTuple): # noqa: D101 - centroids: jnp.ndarray - prev_assignment: jnp.ndarray - assignment: jnp.ndarray - errors: jnp.ndarray + centroids: jax.Array + prev_assignment: jax.Array + assignment: jax.Array + errors: jax.Array center_shift: float class KMeansConst(NamedTuple): # noqa: D101 geom: pointcloud.PointCloud - x_weights: jnp.ndarray + x_weights: jax.Array @property - def x(self) -> jnp.ndarray: + def x(self) -> jax.Array: """Array of shape ``[n, ndim]`` containing the unweighted point cloud.""" return self.geom.x @@ -57,7 +59,7 @@ def weighted_x(self): return self.x_weights[:, :-1] @property - def weights(self) -> jnp.ndarray: + def weights(self) -> jax.Array: """Array of shape ``[n, 1]`` containing weights for each point.""" return self.x_weights[:, -1:] @@ -75,12 +77,12 @@ class KMeansOutput(NamedTuple): inner_errors: Array of shape ``[max_iterations,]`` containing the ``error`` at every iteration. """ - centroids: jnp.ndarray - assignment: jnp.ndarray + centroids: jax.Array + assignment: jax.Array converged: bool iteration: int error: float - inner_errors: Optional[jnp.ndarray] + inner_errors: jax.Array | None @classmethod def _from_state( @@ -110,7 +112,7 @@ def _from_state( def _random_init( geom: pointcloud.PointCloud, k: int, rng: jax.Array -) -> jnp.ndarray: +) -> jax.Array: n, _ = geom.shape ixs = jax.random.choice(rng, jnp.arange(n), shape=(k,), replace=False) return geom.x[ixs] @@ -120,8 +122,8 @@ def _k_means_plus_plus( geom: pointcloud.PointCloud, k: int, rng: jax.Array, - n_local_trials: Optional[int] = None, -) -> jnp.ndarray: + n_local_trials: int | None = None, +) -> jax.Array: def init_fn(geom: pointcloud.PointCloud, rng: jax.Array) -> KPPState: rng, next_rng = jax.random.split(rng, 2) @@ -131,7 +133,7 @@ def init_fn(geom: pointcloud.PointCloud, rng: jax.Array) -> KPPState: return KPPState(rng=next_rng, centroids=centroids, centroid_dists=dists) def body_fn( - iteration: int, const: Tuple[pointcloud.PointCloud, jnp.ndarray], + iteration: int, const: tuple[pointcloud.PointCloud, jax.Array], state: KPPState, compute_error: bool ) -> KPPState: del compute_error @@ -177,10 +179,10 @@ def body_fn( @functools.partial(jax.vmap, in_axes=[None, 0, 0, 0], out_axes=0) def _reallocate_centroids( const: KMeansConst, - ix: jnp.ndarray, - centroid: jnp.ndarray, - weight: jnp.ndarray, -) -> Tuple[jnp.ndarray, jnp.ndarray]: + ix: jax.Array, + centroid: jax.Array, + weight: jax.Array, +) -> tuple[jax.Array, jax.Array]: is_empty = weight <= 0.0 new_centroid = (1 - is_empty) * centroid + is_empty * const.x[ix] # (ndim,) centroid_to_remove = is_empty * const.weighted_x[ix] # (ndim,) @@ -190,8 +192,8 @@ def _reallocate_centroids( def _update_assignment( const: KMeansConst, - centroids: jnp.ndarray, -) -> Tuple[jnp.ndarray, jnp.ndarray]: + centroids: jax.Array, +) -> tuple[jax.Array, jax.Array]: (x, _, *args), aux_data = const.geom.tree_flatten() cost_matrix = type( const.geom @@ -203,9 +205,9 @@ def _update_assignment( def _update_centroids( - const: KMeansConst, k: int, assignment: jnp.ndarray, - dist_to_centers: jnp.ndarray -) -> jnp.ndarray: + const: KMeansConst, k: int, assignment: jax.Array, + dist_to_centers: jax.Array +) -> jax.Array: # TODO(michalk8): # cannot put `k` into `const`, see https://github.com/ott-jax/ott/issues/129 x_weights = jax.ops.segment_sum(const.x_weights, assignment, num_segments=k) @@ -227,9 +229,9 @@ def _k_means( rng: jax.Array, geom: pointcloud.PointCloud, k: int, - weights: Optional[jnp.ndarray] = None, + weights: jax.Array | None = None, init: Init_t = "k-means++", - n_local_trials: Optional[int] = None, + n_local_trials: int | None = None, tol: float = 1e-4, min_iterations: int = 0, max_iterations: int = 300, @@ -342,17 +344,17 @@ def finalize_fn(const: KMeansConst, state: KMeansState) -> KMeansState: def k_means( - geom: Union[jnp.ndarray, pointcloud.PointCloud], + geom: jax.Array | pointcloud.PointCloud, k: int, - weights: Optional[jnp.ndarray] = None, + weights: jax.Array | None = None, init: Init_t = "k-means++", n_init: int = 10, - n_local_trials: Optional[int] = None, + n_local_trials: int | None = None, tol: float = 1e-4, min_iterations: int = 0, max_iterations: int = 300, store_inner_errors: bool = False, - rng: Optional[jax.Array] = None, + rng: jax.Array | None = None, ) -> KMeansOutput: r"""K-means clustering using Lloyd's algorithm :cite:`lloyd:82`. @@ -386,7 +388,7 @@ def k_means( """ assert geom.shape[ 0] >= k, f"Cannot cluster `{geom.shape[0]}` points into `{k}` clusters." - if isinstance(geom, jnp.ndarray): + if isinstance(geom, jax.Array): geom = pointcloud.PointCloud(geom) if isinstance(geom.cost_fn, costs.Cosine): geom = geom._cosine_to_sqeucl() @@ -409,4 +411,4 @@ def k_means( max_iterations, store_inner_errors ) best_ix = jnp.argmin(out.error) - return jax.tree_util.tree_map(lambda arr: arr[best_ix], out) + return jax.tree.map(lambda arr: arr[best_ix], out) diff --git a/src/ott/tools/plot.py b/src/ott/tools/plot.py index 65aa1775f..fb9a58af6 100644 --- a/src/ott/tools/plot.py +++ b/src/ott/tools/plot.py @@ -11,17 +11,8 @@ # 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 typing import ( - Any, - Callable, - Dict, - List, - Literal, - Optional, - Sequence, - Tuple, - Union, -) +from collections.abc import Callable, Sequence +from typing import Any, Literal import jax import jax.numpy as jnp @@ -40,8 +31,10 @@ from ott.solvers.quadratic import gromov_wasserstein # TODO(michalk8): make sure all outputs conform to a unified transport interface -Transport = Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput, - gromov_wasserstein.GWOutput] +Transport = ( + sinkhorn.SinkhornOutput | sinkhorn_lr.LRSinkhornOutput + | gromov_wasserstein.GWOutput +) __all__ = ["Plot", "PlotMM", "transport_animation"] @@ -77,17 +70,17 @@ class Plot: def __init__( self, - fig: Optional["plt.Figure"] = None, - ax: Optional["plt.Axes"] = None, + fig: plt.Figure | None = None, + ax: plt.Axes | None = None, threshold: float = -1.0, scale: int = 200, show_lines: bool = True, cmap: str = "cool", scale_alpha_by_coupling: bool = False, alpha: float = 0.7, - title: Optional[str] = None, - xlim: Optional[List[float]] = None, - ylim: Optional[List[float]] = None, + title: str | None = None, + xlim: list[float] | None = None, + ylim: list[float] | None = None, ): if ax is None and fig is None: fig, ax = plt.subplots() @@ -154,7 +147,7 @@ def _mapping(self, x: jax.Array, y: jax.Array, matrix: jax.Array): return result - def __call__(self, ot: Transport) -> List[plt.Artist]: + def __call__(self, ot: Transport) -> list[plt.Artist]: """Plot couplings in 2-D, using PCA if data is higher dimensional.""" x, y, sx, sy = self._scatter(ot) self._points_x = self.ax.scatter( @@ -192,9 +185,7 @@ def __call__(self, ot: Transport) -> List[plt.Artist]: return [self._points_x, self._points_y] + self._lines - def update(self, - ot: Transport, - title: Optional[str] = None) -> List[plt.Artist]: + def update(self, ot: Transport, title: str | None = None) -> list[plt.Artist]: """Update a plot with a transport instance.""" x, y, _, _ = self._scatter(ot) self._points_x.set_offsets(x) @@ -238,7 +229,7 @@ def update(self, def animate( self, transports: Sequence[Transport], - titles: Optional[Sequence[str]] = None, + titles: Sequence[str] | None = None, frame_rate: float = 10.0 ) -> animation.FuncAnimation: """Make an animation from several transports.""" @@ -282,13 +273,13 @@ class PlotMM(Plot): def __init__( self, - fig: Optional[plt.Figure] = None, - ax: Optional[plt.Axes] = None, + fig: plt.Figure | None = None, + ax: plt.Axes | None = None, fix_axes_lim: bool = False, - cmap: Union[str, mcolors.Colormap] = "cividis_r", + cmap: str | mcolors.Colormap = "cividis_r", markers: str = "svopxdh", alpha: float = 0.6, - title: Optional[str] = None, + title: str | None = None, ): if isinstance(cmap, str): cmap = plt.colormaps[cmap] @@ -298,11 +289,9 @@ def __init__( self._markers = markers self._fix_axes_lim = fix_axes_lim - def __call__( - self, - ot: mmsinkhorn.MMSinkhornOutput, - top_k: Optional[int] = None - ) -> List["plt.Artist"]: + def __call__(self, + ot: mmsinkhorn.MMSinkhornOutput, + top_k: int | None = None) -> list["plt.Artist"]: """Plot 2-D couplings. does not support higher dimensional.""" assert ot.n_marginals <= len(self._markers), "Not enough markers to plot." self._points = [] @@ -352,9 +341,9 @@ def __call__( def update( self, ot: mmsinkhorn.MMSinkhornOutput, - title: Optional[str] = None, - top_k: Optional[int] = None, - ) -> List[plt.Artist]: + title: str | None = None, + top_k: int | None = None, + ) -> list[plt.Artist]: """Update a plot with a transport instance.""" n0 = max(ot.shape) top_k = n0 if top_k is None else top_k @@ -390,9 +379,9 @@ def update( def animate( self, transports: Sequence[mmsinkhorn.MMSinkhornOutput], - titles: Optional[Sequence[str]] = None, + titles: Sequence[str] | None = None, frame_rate: float = 10.0, - top_k: Optional[int] = None, + top_k: int | None = None, ) -> animation.FuncAnimation: """Make an animation from several transports.""" ot, *_ = transports @@ -417,7 +406,7 @@ def get_plotkwargs( small_size: int = 50, mid_size: int = 60, size_multiplier: float = 1.2 -) -> Dict[str, Any]: +) -> dict[str, Any]: r"""Generate marker styling specifications for transport visualization. This utility function creates a dictionary of matplotlib styling parameters @@ -555,20 +544,19 @@ def transport_animation( static_tgt_points: jax.Array, *, n_grid: int = 0, - velocity_field: Optional[Callable[[jax.Array, jax.Array], - jax.Array]] = None, - dynamic_src_points: Optional[jax.Array] = None, - num_ifm_interpolants: Union[int, Literal["all"]] = 0, + velocity_field: Callable[[jax.Array, jax.Array], jax.Array] | None = None, + dynamic_src_points: jax.Array | None = None, + num_ifm_interpolants: int | Literal["all"] = 0, plot_ifm_arrows: bool = False, - title: Optional[str] = None, - figsize: Tuple[int, int] = (8, 6), - xlimits: Optional[Tuple[float, float]] = None, - ylimits: Optional[Tuple[float, float]] = None, + title: str | None = None, + figsize: tuple[int, int] = (8, 6), + xlimits: tuple[float, float] | None = None, + ylimits: tuple[float, float] | None = None, padding: float = 0.1, interval: int = 300, - save_path: Optional[str] = None, + save_path: str | None = None, darkmode: bool = False -) -> Union[plt.Figure, animation.FuncAnimation]: +) -> plt.Figure | animation.FuncAnimation: r"""Create animated visualizations of optimal transport and flow matching. This function generates animations illustrating various aspects of optimal @@ -837,7 +825,7 @@ def ccworder(A: jax.Array) -> jax.Array: return jnp.argsort(jnp.arctan2(A[:, 1], A[:, 0])) -def bidimensional(x: jax.Array, y: jax.Array) -> Tuple[jax.Array, jax.Array]: +def bidimensional(x: jax.Array, y: jax.Array) -> tuple[jax.Array, jax.Array]: """Apply PCA to reduce to bi-dimensional data.""" if x.shape[1] < 3: return x, y diff --git a/src/ott/tools/progot.py b/src/ott/tools/progot.py index 0fd273b9f..ff7b84506 100644 --- a/src/ott/tools/progot.py +++ b/src/ott/tools/progot.py @@ -11,7 +11,7 @@ # 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 typing import Any, Literal, NamedTuple, Optional, Tuple, Union +from typing import Any, Literal, NamedTuple import jax import jax.numpy as jnp @@ -30,12 +30,12 @@ "get_alpha_schedule", ] -Output = Union[sinkhorn.SinkhornOutput, sd.SinkhornDivergenceOutput] +Output = sinkhorn.SinkhornOutput | sd.SinkhornDivergenceOutput class ProgOTState(NamedTuple): - x: jnp.ndarray - init_potentials: Optional[Tuple[jnp.ndarray, jnp.ndarray]] + x: jax.Array + init_potentials: tuple[jax.Array, jax.Array] | None class ProgOTOutput(NamedTuple): @@ -49,17 +49,17 @@ class ProgOTOutput(NamedTuple): xs: Intermediate interpolations of shape ``[num_steps, n, d]``, if present. """ prob: linear_problem.LinearProblem - alphas: jnp.ndarray - epsilons: jnp.ndarray + alphas: jax.Array + epsilons: jax.Array outputs: Output - xs: Optional[jnp.ndarray] = None + xs: jax.Array | None = None def transport( self, - x: jnp.ndarray, - num_steps: Optional[int] = None, + x: jax.Array, + num_steps: int | None = None, return_intermediate: bool = False, - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + ) -> tuple[jax.Array, jax.Array]: """Transport points. Args: @@ -76,9 +76,9 @@ def transport( """ def body_fn( - xy: Tuple[jnp.ndarray, Optional[jnp.ndarray]], it: int - ) -> Tuple[Tuple[jnp.ndarray, Optional[jnp.ndarray]], Tuple[ - Optional[jnp.ndarray], Optional[jnp.ndarray]]]: + xy: tuple[jax.Array, jax.Array | None], it: int + ) -> tuple[tuple[jax.Array, jax.Array | None], tuple[jax.Array | None, + jax.Array | None]]: x, _ = xy alpha = self.alphas[it] dp = self.get_output(it).to_dual_potentials() @@ -111,12 +111,10 @@ def get_output(self, step: int) -> Output: Returns: The OT solver output at a ``step``. """ - return jtu.tree_map(lambda x: x[step], self.outputs) + return jax.tree.map(lambda x: x[step], self.outputs) @property - def converged( - self - ) -> Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]]: + def converged(self) -> jax.Array | tuple[jax.Array, jax.Array, jax.Array]: """Convergence at each step. - If :attr:`is_debiased`, return an array of shape ``[num_steps, 3]`` with @@ -127,7 +125,7 @@ def converged( return jnp.stack(self.outputs.converged, axis=-1) @property - def num_iters(self) -> jnp.ndarray: + def num_iters(self) -> jax.Array: """Number of Sinkhorn iterations within each step. - If :attr:`is_debiased`, return an array of shape ``[num_steps, 3]`` with @@ -168,10 +166,10 @@ class ProgOT: def __init__( self, - alphas: jnp.ndarray, + alphas: jax.Array, *, - epsilons: Optional[jnp.ndarray] = None, - epsilon_scales: Optional[jnp.ndarray] = None, + epsilons: jax.Array | None = None, + epsilon_scales: jax.Array | None = None, is_debiased: bool = False, ): if epsilons is not None and epsilon_scales is not None: @@ -218,7 +216,7 @@ def __call__( """ def body_fn(state: ProgOTState, - it: int) -> Tuple[ProgOTState, Tuple[Output, float]]: + it: int) -> tuple[ProgOTState, tuple[Output, float]]: alpha = self.alphas[it] eps = None if self.epsilons is None else self.epsilons[it] if self.epsilon_scales is not None: @@ -298,12 +296,12 @@ def tree_unflatten( # noqa: D102 def get_epsilon_schedule( geom: pointcloud.PointCloud, *, - alphas: jnp.ndarray, - epsilon_scales: jnp.ndarray, - y_eval: jnp.ndarray, + alphas: jax.Array, + epsilon_scales: jax.Array, + y_eval: jax.Array, start_epsilon_scale: float = 1.0, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: """Get the epsilon regularization schedule. See Algorithm 4 in :cite:`kassraie:24` for more information. @@ -358,7 +356,7 @@ def error(epsilon_scale: float) -> float: def get_alpha_schedule( kind: Literal["lin", "exp", "quad"], *, num_steps: int -) -> jnp.ndarray: +) -> jax.Array: """Get the step size schedule. Convenience wrapper to get a sequence of ``num_steps`` timestamps between @@ -391,11 +389,11 @@ def get_alpha_schedule( def _sinkhorn( - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, cost_fn: costs.TICost, - eps: Optional[float], - init: Optional[Tuple[jnp.ndarray, jnp.ndarray]] = None, + eps: float | None, + init: tuple[jax.Array, jax.Array] | None = None, **kwargs: Any, ) -> sinkhorn.SinkhornOutput: geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn, epsilon=eps) @@ -405,10 +403,10 @@ def _sinkhorn( def _sinkhorn_divergence( - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, cost_fn: costs.TICost, - eps: Optional[float], + eps: float | None, **kwargs: Any, ) -> sd.SinkhornDivergenceOutput: _, out = sd.sinkhorn_divergence( diff --git a/src/ott/tools/segment_sinkhorn.py b/src/ott/tools/segment_sinkhorn.py index 52f3295b6..62c6ae9c7 100644 --- a/src/ott/tools/segment_sinkhorn.py +++ b/src/ott/tools/segment_sinkhorn.py @@ -11,10 +11,11 @@ # 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 collections.abc import Mapping from types import MappingProxyType -from typing import Any, Mapping, Optional, Tuple +from typing import Any -import jax.numpy as jnp +import jax from ott.geometry import costs, pointcloud, segment from ott.problems.linear import linear_problem @@ -22,21 +23,21 @@ def segment_sinkhorn( - x: jnp.ndarray, - y: jnp.ndarray, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, + x: jax.Array, + y: jax.Array, + num_segments: int | None = None, + max_measure_size: int | None = None, + cost_fn: costs.CostFn | None = None, + segment_ids_x: jax.Array | None = None, + segment_ids_y: jax.Array | None = None, indices_are_sorted: bool = False, - num_per_segment_x: Optional[Tuple[int, ...]] = None, - num_per_segment_y: Optional[Tuple[int, ...]] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, + num_per_segment_x: tuple[int, ...] | None = None, + num_per_segment_y: tuple[int, ...] | None = None, + weights_x: jax.Array | None = None, + weights_y: jax.Array | None = None, sinkhorn_kwargs: Mapping[str, Any] = MappingProxyType({}), **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Compute regularized OT cost between subsets of vectors in ``x`` and ``y``. Helper function designed to compute Sinkhorn regularized OT cost between @@ -104,11 +105,11 @@ def segment_sinkhorn( padding_vector = cost_fn._padder(dim=dim) def eval_fn( - padded_x: jnp.ndarray, - padded_y: jnp.ndarray, - padded_weight_x: jnp.ndarray, - padded_weight_y: jnp.ndarray, - ) -> jnp.ndarray: + padded_x: jax.Array, + padded_y: jax.Array, + padded_weight_x: jax.Array, + padded_weight_y: jax.Array, + ) -> jax.Array: geom = pointcloud.PointCloud( padded_x, padded_y, diff --git a/src/ott/tools/sinkhorn_divergence.py b/src/ott/tools/sinkhorn_divergence.py index 5dd56b4c7..168ccbdca 100644 --- a/src/ott/tools/sinkhorn_divergence.py +++ b/src/ott/tools/sinkhorn_divergence.py @@ -11,8 +11,9 @@ # 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 collections.abc import Mapping from types import MappingProxyType -from typing import Any, Mapping, Optional, Tuple, Type, Union +from typing import Any import jax import jax.numpy as jnp @@ -28,8 +29,8 @@ "SinkhornDivergenceOutput" ] -Potentials = Tuple[jnp.ndarray, jnp.ndarray] -Factors = Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] +Potentials = tuple[jax.Array, jax.Array] +Factors = tuple[jax.Array, jax.Array, jax.Array] @utils.register_pytree_node @@ -63,18 +64,17 @@ class SinkhornDivergenceOutput: # noqa: D101 complete each of the three terms in the divergence. """ divergence: float - geoms: Tuple[geometry.Geometry, geometry.Geometry, geometry.Geometry] - a: jnp.ndarray - b: jnp.ndarray - potentials: Optional[Tuple[Potentials, Potentials, Potentials]] - factors: Optional[Tuple[Factors, Factors, Factors]] - errors: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray], - Optional[jnp.ndarray]] - converged: Tuple[bool, bool, bool] - n_iters: Tuple[int, int, int] + geoms: tuple[geometry.Geometry, geometry.Geometry, geometry.Geometry] + a: jax.Array + b: jax.Array + potentials: tuple[Potentials, Potentials, Potentials] | None + factors: tuple[Factors, Factors, Factors] | None + errors: tuple[jax.Array | None, jax.Array | None, jax.Array | None] + converged: tuple[bool, bool, bool] + n_iters: tuple[int, int, int] def to_dual_potentials( - self, epsilon: Optional[float] = None + self, epsilon: float | None = None ) -> potentials.DualPotentials: """Return dual potential functions :cite:`pooladian:22`. @@ -134,13 +134,13 @@ def tree_unflatten(cls, aux_data, children): # noqa: D102 def sinkdiv( - x: jnp.ndarray, - y: jnp.ndarray, + x: jax.Array, + y: jax.Array, *, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, **kwargs: Any, -) -> Tuple[jnp.ndarray, SinkhornDivergenceOutput]: +) -> tuple[jax.Array, SinkhornDivergenceOutput]: """Wrapper to get the :term:`Sinkhorn divergence` between two point clouds. Convenience wrapper around @@ -176,17 +176,17 @@ def sinkdiv( def sinkhorn_divergence( - geom: Type[geometry.Geometry], + geom: type[geometry.Geometry], *args: Any, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, + a: jax.Array | None = None, + b: jax.Array | None = None, solve_kwargs: Mapping[str, Any] = MappingProxyType({}), static_b: bool = False, - offset_static_b: Optional[float] = None, + offset_static_b: float | None = None, share_epsilon: bool = True, symmetric_sinkhorn: bool = True, **kwargs: Any, -) -> Tuple[jnp.ndarray, SinkhornDivergenceOutput]: +) -> tuple[jax.Array, SinkhornDivergenceOutput]: r"""Compute :term:`Sinkhorn divergence` between two measures. The :term:`Sinkhorn divergence` is computed between two measures :math:`\mu` @@ -258,11 +258,11 @@ def sinkhorn_divergence( def _sinkhorn_divergence( geometry_xy: geometry.Geometry, geometry_xx: geometry.Geometry, - geometry_yy: Optional[geometry.Geometry], - a: jnp.ndarray, - b: jnp.ndarray, + geometry_yy: geometry.Geometry | None, + a: jax.Array, + b: jax.Array, symmetric_sinkhorn: bool, - offset_yy: Optional[float], + offset_yy: float | None, **kwargs: Any, ) -> SinkhornDivergenceOutput: """Compute the (unbalanced) Sinkhorn divergence for the wrapper function. @@ -279,9 +279,9 @@ def _sinkhorn_divergence( between elements of the view X. geometry_yy: a Cost object able to apply kernels with a certain epsilon, between elements of the view Y. - a: jnp.ndarray[n]: the weight of each input point. The sum of + a: jax.Array[n]: the weight of each input point. The sum of all elements of ``b`` must match that of ``a`` to converge. - b: jnp.ndarray[m]: the weight of each target point. The sum of + b: jax.Array[m]: the weight of each target point. The sum of all elements of ``b`` must match that of ``a`` to converge. symmetric_sinkhorn: Use Sinkhorn updates in Eq. 25 of :cite:`feydy:19` for symmetric terms comparing x/x and y/y. @@ -352,24 +352,24 @@ def _sinkhorn_divergence( def segment_sinkhorn_divergence( - x: jnp.ndarray, - y: jnp.ndarray, - num_segments: Optional[int] = None, - max_measure_size: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - segment_ids_x: Optional[jnp.ndarray] = None, - segment_ids_y: Optional[jnp.ndarray] = None, + x: jax.Array, + y: jax.Array, + num_segments: int | None = None, + max_measure_size: int | None = None, + cost_fn: costs.CostFn | None = None, + segment_ids_x: jax.Array | None = None, + segment_ids_y: jax.Array | None = None, indices_are_sorted: bool = False, - num_per_segment_x: Optional[Tuple[int, ...]] = None, - num_per_segment_y: Optional[Tuple[int, ...]] = None, - weights_x: Optional[jnp.ndarray] = None, - weights_y: Optional[jnp.ndarray] = None, + num_per_segment_x: tuple[int, ...] | None = None, + num_per_segment_y: tuple[int, ...] | None = None, + weights_x: jax.Array | None = None, + weights_y: jax.Array | None = None, solve_kwargs: Mapping[str, Any] = MappingProxyType({}), static_b: bool = False, share_epsilon: bool = True, symmetric_sinkhorn: bool = False, **kwargs: Any -) -> jnp.ndarray: +) -> jax.Array: """Compute Sinkhorn divergence between subsets of vectors in `x` and `y`. Helper function designed to compute Sinkhorn divergences between several point @@ -446,10 +446,10 @@ def segment_sinkhorn_divergence( padding_vector = cost_fn._padder(dim=dim) def eval_fn( - padded_x: jnp.ndarray, - padded_y: jnp.ndarray, - padded_weight_x: jnp.ndarray, - padded_weight_y: jnp.ndarray, + padded_x: jax.Array, + padded_y: jax.Array, + padded_weight_x: jax.Array, + padded_weight_y: jax.Array, ) -> float: div, _ = sinkhorn_divergence( pointcloud.PointCloud, @@ -485,8 +485,8 @@ def eval_fn( def _empty_output( is_low_rank: bool, - offset_yy: Optional[float] = None -) -> Union[sinkhorn.SinkhornOutput, sinkhorn_lr.LRSinkhornOutput]: + offset_yy: float | None = None +) -> sinkhorn.SinkhornOutput | sinkhorn_lr.LRSinkhornOutput: if is_low_rank: return sinkhorn_lr.LRSinkhornOutput( q=None, diff --git a/src/ott/tools/sliced.py b/src/ott/tools/sliced.py index a37dc6449..02b2cf22b 100644 --- a/src/ott/tools/sliced.py +++ b/src/ott/tools/sliced.py @@ -11,7 +11,7 @@ # 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 typing import Callable, Optional, Tuple +from collections.abc import Callable import jax import jax.numpy as jnp @@ -23,15 +23,15 @@ __all__ = ["random_proj_sphere", "sliced_wasserstein"] -Projector = Callable[[jax.Array, jnp.ndarray], jnp.ndarray] +Projector = Callable[[jax.Array, jax.Array], jax.Array] def random_proj_sphere( rng: jax.Array, - x: jnp.ndarray, + x: jax.Array, *, n_proj: int = 1000, -) -> jnp.ndarray: +) -> jax.Array: """Project data on directions sampled randomly from sphere. Args: @@ -50,17 +50,17 @@ def random_proj_sphere( def sliced_wasserstein( - x: jnp.ndarray, - y: jnp.ndarray, - a: Optional[jnp.ndarray] = None, - b: Optional[jnp.ndarray] = None, - cost_fn: Optional[costs.CostFn] = None, - proj_fn: Optional[Projector] = None, - weights: Optional[jnp.ndarray] = None, + x: jax.Array, + y: jax.Array, + a: jax.Array | None = None, + b: jax.Array | None = None, + cost_fn: costs.CostFn | None = None, + proj_fn: Projector | None = None, + weights: jax.Array | None = None, return_transport: bool = False, return_dual_variables: bool = False, - rng: Optional[jax.Array] = None, -) -> Tuple[jnp.ndarray, univariate.UnivariateOutput]: + rng: jax.Array | None = None, +) -> tuple[jax.Array, univariate.UnivariateOutput]: r"""Compute the Sliced Wasserstein distance between two weighted point clouds. Follows the approach outlined in :cite:`rabin:12` to compute a proxy for OT diff --git a/src/ott/tools/soft_sort.py b/src/ott/tools/soft_sort.py index 977cd333f..c4f65f51c 100644 --- a/src/ott/tools/soft_sort.py +++ b/src/ott/tools/soft_sort.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Callable, Optional, Tuple, Union +from collections.abc import Callable +from typing import Any import jax import jax.numpy as jnp import numpy as np +from jax.typing import ArrayLike from ott import utils from ott.geometry import costs, pointcloud @@ -30,14 +32,14 @@ "quantize", "topk_mask", "multivariate_cdf_quantile_maps" ] -Func_t = Callable[[jnp.ndarray], jnp.ndarray] +Func_t = Callable[[jax.Array], jax.Array] def transport_for_sort( - inputs: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, - target_weights: Optional[jnp.ndarray] = None, - squashing_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + inputs: jax.Array, + weights: jax.Array | None = None, + target_weights: jax.Array | None = None, + squashing_fun: Callable[[jax.Array], jax.Array] | None = None, epsilon: float = 1e-2, **kwargs: Any, ) -> sinkhorn.SinkhornOutput: @@ -83,7 +85,7 @@ def transport_for_sort( return solver(prob) -def apply_on_axis(op, inputs, axis, *args, **kwargs: Any) -> jnp.ndarray: +def apply_on_axis(op, inputs, axis, *args, **kwargs: Any) -> jax.Array: """Apply a differentiable operator on a given axis of the input. Args: @@ -120,8 +122,8 @@ def apply_on_axis(op, inputs, axis, *args, **kwargs: Any) -> jnp.ndarray: def _sort( - inputs: jnp.ndarray, topk: int, num_targets: Optional[int], **kwargs: Any -) -> jnp.ndarray: + inputs: jax.Array, topk: int, num_targets: int | None, **kwargs: Any +) -> jax.Array: """Apply the soft sort operator on a one dimensional array.""" num_points = inputs.shape[0] a = jnp.ones((num_points,)) / num_points @@ -145,12 +147,12 @@ def _sort( def sort( - inputs: jnp.ndarray, + inputs: jax.Array, axis: int = -1, topk: int = -1, - num_targets: Optional[int] = None, + num_targets: int | None = None, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Apply the soft sort operator on a given axis of the input. For instance: @@ -203,8 +205,8 @@ def sort( def _ranks( - inputs: jnp.ndarray, num_targets, target_weights, **kwargs: Any -) -> jnp.ndarray: + inputs: jax.Array, num_targets, target_weights, **kwargs: Any +) -> jax.Array: """Apply the soft ranks operator on a one dimensional array.""" num_points = inputs.shape[0] if target_weights is None: @@ -220,12 +222,12 @@ def _ranks( def ranks( - inputs: jnp.ndarray, + inputs: jax.Array, axis: int = -1, - num_targets: Optional[int] = None, - target_weights: Optional[jnp.ndarray] = None, + num_targets: int | None = None, + target_weights: jax.Array | None = None, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Apply the soft rank operator on input tensor. For instance: @@ -278,11 +280,11 @@ def ranks( def topk_mask( - inputs: jnp.ndarray, + inputs: jax.Array, axis: int = -1, k: int = 1, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Soft :math:`\text{top-}k` selection mask. For instance: @@ -337,12 +339,12 @@ def topk_mask( def quantile( - inputs: jnp.ndarray, - q: Optional[Union[float, jnp.ndarray]], - axis: Union[int, Tuple[int, ...]] = -1, - weight: Optional[Union[float, jnp.ndarray]] = None, + inputs: jax.Array, + q: ArrayLike | None, + axis: int | tuple[int, ...] = -1, + weight: ArrayLike | None = None, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Apply the soft quantiles operator on the input tensor. For instance: @@ -396,8 +398,8 @@ def quantile( """ def _quantile( - inputs: jnp.ndarray, q: float, weight: float, **kwargs - ) -> jnp.ndarray: + inputs: jax.Array, q: float, weight: float, **kwargs + ) -> jax.Array: num_points = inputs.shape[0] q = jnp.array([0.2, 0.5, 0.8]) if q is None else jnp.atleast_1d(q) num_quantiles = q.shape[0] @@ -457,15 +459,15 @@ def _quantile( def multivariate_cdf_quantile_maps( - inputs: jnp.ndarray, - target_sampler: Optional[Callable[[jax.Array, Tuple[int, int]], - jax.Array]] = None, - rng: Optional[jax.Array] = None, - num_target_samples: Optional[int] = None, - cost_fn: Optional[costs.CostFn] = None, - epsilon: Optional[float] = None, - input_weights: Optional[jnp.ndarray] = None, - target_weights: Optional[jnp.ndarray] = None, + inputs: jax.Array, + target_sampler: Callable[[jax.Array, tuple[int, int]], jax.Array] + | None = None, + rng: jax.Array | None = None, + num_target_samples: int | None = None, + cost_fn: costs.CostFn | None = None, + epsilon: float | None = None, + input_weights: jax.Array | None = None, + target_weights: jax.Array | None = None, **kwargs: Any ) -> potentials.DualPotentials: r"""Returns multivariate CDF and quantile maps, given input samples. @@ -535,8 +537,8 @@ def multivariate_cdf_quantile_maps( def _quantile_normalization( - inputs: jnp.ndarray, targets: jnp.ndarray, weights: float, **kwargs: Any -) -> jnp.ndarray: + inputs: jax.Array, targets: jax.Array, weights: float, **kwargs: Any +) -> jax.Array: """Apply soft quantile normalization on a one dimensional array.""" num_points = inputs.shape[0] a = jnp.ones((num_points,)) / num_points @@ -545,12 +547,12 @@ def _quantile_normalization( def quantile_normalization( - inputs: jnp.ndarray, - targets: jnp.ndarray, - weights: Optional[jnp.ndarray] = None, + inputs: jax.Array, + targets: jax.Array, + weights: jax.Array | None = None, axis: int = -1, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Re-normalize inputs so that its quantiles match those of targets/weights. Quantile normalization rearranges the values in inputs to values that match @@ -601,11 +603,11 @@ def quantile_normalization( def sort_with( - inputs: jnp.ndarray, - criterion: jnp.ndarray, + inputs: jax.Array, + criterion: jax.Array, topk: int = -1, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Sort a multidimensional array according to a real valued criterion. Given ``batch`` vectors of dimension `dim`, to which, for each, a real value @@ -656,7 +658,7 @@ def sort_with( return sort_fn(inputs) -def _quantize(inputs: jnp.ndarray, num_q: int, **kwargs: Any) -> jnp.ndarray: +def _quantize(inputs: jax.Array, num_q: int, **kwargs: Any) -> jax.Array: """Apply the soft quantization operator on a one dimensional array.""" num_points = inputs.shape[0] a = jnp.ones((num_points,)) / num_points @@ -666,11 +668,11 @@ def _quantize(inputs: jnp.ndarray, num_q: int, **kwargs: Any) -> jnp.ndarray: def quantize( - inputs: jnp.ndarray, + inputs: jax.Array, num_levels: int = 10, axis: int = -1, **kwargs: Any, -) -> jnp.ndarray: +) -> jax.Array: r"""Soft quantizes an input according using ``num_levels`` values along axis. The quantization operator consists in concentrating several values around diff --git a/src/ott/tools/unreg.py b/src/ott/tools/unreg.py index 260c43a16..84695e5a1 100644 --- a/src/ott/tools/unreg.py +++ b/src/ott/tools/unreg.py @@ -11,8 +11,8 @@ # 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 typing import Tuple +import jax import jax.numpy as jnp from optax import assignment @@ -26,7 +26,7 @@ def hungarian( geom: geometry.Geometry -) -> Tuple[jnp.ndarray, semidiscrete.HardAssignmentOutput]: +) -> tuple[jax.Array, semidiscrete.HardAssignmentOutput]: """Solve matching problem using the :term:`Hungarian algorithm`. Uses the implementation from :mod:`optax`. @@ -51,7 +51,7 @@ def hungarian( return transport_cost, out -def wassdis_p(x: jnp.ndarray, y: jnp.ndarray, *, p: float = 2.0) -> float: +def wassdis_p(x: jax.Array, y: jax.Array, *, p: float = 2.0) -> float: """Compute the :term:`Wasserstein distance`, uses :term:`Hungarian algorithm`. Uses :func:`hungarian` to solve the :term:`optimal matching problem` between diff --git a/src/ott/types.py b/src/ott/types.py index e76af71d7..8f2a17dec 100644 --- a/src/ott/types.py +++ b/src/ott/types.py @@ -13,7 +13,7 @@ # limitations under the License. from typing import Protocol -import jax.numpy as jnp +import jax __all__ = ["Transport"] @@ -28,11 +28,11 @@ class can however be used in type hints to support duck typing. """ @property - def matrix(self) -> jnp.ndarray: + def matrix(self) -> jax.Array: ... - def apply(self, inputs: jnp.ndarray, axis: int) -> jnp.ndarray: + def apply(self, inputs: jax.Array, axis: int) -> jax.Array: ... - def marginal(self, axis: int = 0) -> jnp.ndarray: + def marginal(self, axis: int = 0) -> jax.Array: ... diff --git a/src/ott/utils.py b/src/ott/utils.py index fc1166d59..8d428059d 100644 --- a/src/ott/utils.py +++ b/src/ott/utils.py @@ -15,18 +15,8 @@ import functools import io import warnings -from collections.abc import Sequence -from typing import ( - Any, - Callable, - List, - NamedTuple, - Optional, - ParamSpec, - Tuple, - TypeVar, - Union, -) +from collections.abc import Callable, Sequence +from typing import Any, NamedTuple, ParamSpec, TypeVar import jax import jax._src.interpreters.batching as batching @@ -47,7 +37,7 @@ "batched_vmap", ] -IOStatus = Tuple[np.ndarray, np.ndarray, np.ndarray, NamedTuple] +IOStatus = tuple[np.ndarray, np.ndarray, np.ndarray, NamedTuple] IOCallback = Callable[[IOStatus], None] P = ParamSpec("P") R = TypeVar("R") @@ -56,7 +46,7 @@ def register_pytree_node(cls: type) -> type: """Register dataclasses as pytree_nodes.""" cls = dataclasses.dataclass()(cls) - flatten = lambda obj: jax.tree_util.tree_flatten(dataclasses.asdict(obj)) + flatten = lambda obj: jax.tree.flatten(dataclasses.asdict(obj)) unflatten = lambda d, children: cls(**d.unflatten(children)) jax.tree_util.register_pytree_node(cls, flatten, unflatten) return cls @@ -64,9 +54,9 @@ def register_pytree_node(cls: type) -> type: def deprecate( # noqa: D103 *, - version: Optional[str] = None, - alt: Optional[str] = None, - func: Optional[Callable[P, R]] = None + version: str | None = None, + alt: str | None = None, + func: Callable[P, R] | None = None ) -> Callable[P, R]: def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: @@ -84,7 +74,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: return functools.wraps(func)(wrapper) -def default_prng_key(rng: Optional[jax.Array] = None) -> jax.Array: +def default_prng_key(rng: jax.Array | None = None) -> jax.Array: """Get the default PRNG key. Args: @@ -99,7 +89,7 @@ def default_prng_key(rng: Optional[jax.Array] = None) -> jax.Array: def default_progress_fn( fmt: str = "{iter} / {max_iter} -- {error}", - stream: Optional[io.TextIOBase] = None, + stream: io.TextIOBase | None = None, ) -> IOCallback: """Return a callback that prints the progress when solving :mod:`linear problems `. @@ -216,7 +206,7 @@ def progress_callback(status: IOStatus) -> None: return progress_callback -def _prepare_info(status: IOStatus) -> Tuple[int, int, int, np.ndarray]: +def _prepare_info(status: IOStatus) -> tuple[int, int, int, np.ndarray]: iteration, inner_iterations, total_iter, state = status iteration = int(iteration) + 1 inner_iterations = int(inner_iterations) @@ -260,8 +250,8 @@ def _batch_and_remainder( args: Any, *, batch_size: int, - in_axes: Optional[Union[int, Sequence[int], Any]], -) -> Tuple[Any, Any]: + in_axes: int | Sequence[int] | Any | None, +) -> tuple[Any, Any]: assert batch_size > 0, f"Batch size must be positive, got {batch_size}." leaves, treedef = jax.tree.flatten(args, is_leaf=batching.is_vmappable) in_axes = _prepare_axes( @@ -311,11 +301,10 @@ def _batch_and_remainder( def _apply_scan( - vmapped_fun: Callable[P, R], in_axes: Optional[Union[int, Sequence[int], - Any]] + vmapped_fun: Callable[P, R], in_axes: int | Sequence[int] | Any | None ) -> Callable[P, R]: - def num_steps(axes: List[Any], args: Tuple[Any, ...]) -> int: + def num_steps(axes: list[Any], args: tuple[Any, ...]) -> int: ix = next(ix for ix, axis in enumerate(axes) if axis is not None) leaf, *_ = jax.tree.leaves(args[ix]) axis, *_ = jax.tree.leaves(axes[ix]) @@ -329,7 +318,7 @@ def select_batch(arg: Any, index: int, *, axis: int) -> Any: def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: - def body_fn(carry: None, index: int) -> Tuple[None, R]: + def body_fn(carry: None, index: int) -> tuple[None, R]: del carry new_args = treedef.unflatten([ arg if axis is None else select_batch(arg, index, axis=axis) @@ -358,7 +347,7 @@ def batched_vmap( fun: Callable[P, R], *, batch_size: int, - in_axes: Optional[Union[int, Sequence[int], Any]] = 0, + in_axes: int | Sequence[int] | Any | None = 0, out_axes: Any = 0, ) -> Callable[P, R]: """Batched version of :func:`~jax.vmap`. @@ -383,15 +372,15 @@ def _batched_map( fun: Callable[P, R], *, batch_size: int, - in_axes: Optional[Union[int, Sequence[int], Any]] = 0, + in_axes: int | Sequence[int] | Any | None = 0, out_axes: Any = 0, ) -> Callable[P, R]: - def unbatch(axis: int, x: jnp.ndarray) -> jnp.ndarray: + def unbatch(axis: int, x: jax.Array) -> jax.Array: x = jnp.moveaxis(x, 0, axis) return jax.lax.collapse(x, axis, axis + 2) - def concat(axis: int, x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + def concat(axis: int, x: jax.Array, y: jax.Array) -> jax.Array: return jnp.concatenate([x, y], axis=axis) @functools.wraps(fun) diff --git a/tests/_utils.py b/tests/_utils.py index 4962580c6..b1c785106 100644 --- a/tests/_utils.py +++ b/tests/_utils.py @@ -13,7 +13,8 @@ # limitations under the License. """Helpers shared across the test suite.""" import dataclasses -from typing import Any, Callable, Dict, List, Sequence, Tuple +from collections.abc import Callable, Sequence +from typing import Any import jax import jax.numpy as jnp @@ -32,10 +33,10 @@ @dataclasses.dataclass(frozen=True) class PointClouds: """Two weighted point clouds, the ingredients of a linear OT problem.""" - x: jnp.ndarray # (n, dim) - y: jnp.ndarray # (m, dim) - a: jnp.ndarray # (n,), on the simplex - b: jnp.ndarray # (m,), on the simplex + x: jax.Array # (n, dim) + y: jax.Array # (m, dim) + a: jax.Array # (n,), on the simplex + b: jax.Array # (m,), on the simplex @property def n(self) -> int: @@ -69,8 +70,8 @@ def problem( @dataclasses.dataclass(frozen=True) class QuadClouds(PointClouds): """Two weighted clouds plus the intra-domain costs of a quadratic problem.""" - cx: jnp.ndarray # (n, n) - cy: jnp.ndarray # (m, m) + cx: jax.Array # (n, n) + cy: jax.Array # (m, m) def quad_problem( self, @@ -96,7 +97,7 @@ def random_weights( *, offset: float = 0.1, zero_at: Sequence[int] = (), -) -> jnp.ndarray: +) -> jax.Array: """Sample ``n`` random weights on the simplex, 0 at ``zero_at``.""" a = jr.uniform(rng, (n,)) + offset a = a.at[jnp.asarray(zero_at, dtype=int)].set(0.0) @@ -123,22 +124,22 @@ def random_clouds( ) -def proj(matrix: jnp.ndarray, nu: float = 1.0) -> jnp.ndarray: +def proj(matrix: jax.Array, nu: float = 1.0) -> jax.Array: """Project a matrix onto the Stiefel manifold, scaled by ``nu``.""" assert nu > 0.0, nu u, _, v_h = jnp.linalg.svd(matrix, full_matrices=False) return u.dot(v_h) * jnp.sqrt(nu) -def tracing_progress_fn() -> Tuple[Dict[str, List[Any]], Callable[..., None]]: +def tracing_progress_fn() -> tuple[dict[str, list[Any]], Callable[..., None]]: """Progress callback recording the iterations at which it was called. Returns: The recorded values and the callback to pass as ``progress_fn``. """ - traced: Dict[str, List[Any]] = {"iters": [], "error": [], "total": []} + traced: dict[str, list[Any]] = {"iters": [], "error": [], "total": []} - def progress_fn(status: Tuple[Any, ...], *args: Any) -> None: + def progress_fn(status: tuple[Any, ...], *args: Any) -> None: iteration, inner_iterations, total_iter, state = status iteration = int(iteration) inner_iterations = int(inner_iterations) diff --git a/tests/conftest.py b/tests/conftest.py index 5d8ad324d..6e8ba9b53 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -13,7 +13,8 @@ # limitations under the License. import itertools from collections import abc -from typing import Any, Iterator, Mapping, Optional, Sequence +from collections.abc import Iterator, Mapping, Sequence +from typing import Any import pytest @@ -51,8 +52,8 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: fast_marks = [m for m in metafunc.function.pytestmark if m.name == "fast"] if fast_marks: mark, = fast_marks - selected: Optional[Mapping[str, Any]] = mark.kwargs.pop("only_fast", None) - ids: Optional[Sequence[str]] = mark.kwargs.pop("ids", None) + selected: Mapping[str, Any] | None = mark.kwargs.pop("only_fast", None) + ids: Sequence[str] | None = mark.kwargs.pop("ids", None) if mark.args: argnames, argvalues = mark.args diff --git a/tests/geometry/_graphs.py b/tests/geometry/_graphs.py index bf33a2a38..37062dc9c 100644 --- a/tests/geometry/_graphs.py +++ b/tests/geometry/_graphs.py @@ -13,12 +13,12 @@ # limitations under the License. """Random graphs shared by the graph and geodesic geometry tests.""" import functools -from typing import Optional, Union import networkx as nx from networkx.algorithms import shortest_paths from networkx.generators import random_graphs +import jax import jax.experimental.sparse as jesp import jax.numpy as jnp import numpy as np @@ -28,16 +28,16 @@ __all__ = ["random_graph", "gt_geometry"] -@functools.lru_cache(maxsize=None) +@functools.cache def random_graph( n: int, p: float = 0.3, - seed: Optional[int] = 0, + seed: int | None = 0, is_sparse: bool = False, *, return_laplacian: bool = False, directed: bool = False, -) -> jnp.ndarray: +) -> jax.Array: """Sample a connected random graph with uniformly weighted edges.""" G = random_graphs.fast_gnp_random_graph(n, p, seed=seed, directed=directed) if not directed: @@ -57,9 +57,7 @@ def random_graph( def gt_geometry( - G: Union[jnp.ndarray, nx.Graph], - *, - epsilon: float = 1e-2 + G: jax.Array | nx.Graph, *, epsilon: float = 1e-2 ) -> geometry.Geometry: """Geometry whose cost is the squared shortest-path distance on ``G``.""" if not isinstance(G, nx.Graph): diff --git a/tests/geometry/costs_test.py b/tests/geometry/costs_test.py index 6e285af71..07b356747 100644 --- a/tests/geometry/costs_test.py +++ b/tests/geometry/costs_test.py @@ -344,7 +344,7 @@ def test_reg_transport_fn( @jax.jit @functools.partial(jax.vmap, in_axes=0) - def expected_fn(x: jnp.ndarray) -> jnp.ndarray: + def expected_fn(x: jax.Array) -> jax.Array: f_h = cost_fn.h_transform(f) return x - cost_fn.regularizer.prox(jax.grad(f_h)(x)) diff --git a/tests/geometry/geodesic_test.py b/tests/geometry/geodesic_test.py index 05bb3e04c..f1a32b88f 100644 --- a/tests/geometry/geodesic_test.py +++ b/tests/geometry/geodesic_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional import networkx as nx from networkx.generators import balanced_tree @@ -30,7 +29,7 @@ from tests.geometry import _graphs -def exact_heat_kernel(G: jnp.ndarray, normalize: bool = False, t: float = 10): +def exact_heat_kernel(G: jax.Array, normalize: bool = False, t: float = 10): degree = jnp.sum(G, axis=1) L = jnp.diag(degree) - G if normalize: @@ -88,7 +87,7 @@ def test_kernel_is_symmetric_positive_definite( t=[1e-4, 1e-5], only_fast=-1, ) - def test_approximates_ground_truth(self, t: Optional[float], order: int): + def test_approximates_ground_truth(self, t: float | None, order: int): tol = 1e-2 G = nx.linalg.adjacency_matrix(balanced_tree(r=2, h=5)) G = jnp.asarray(G.toarray().astype(float)) @@ -106,7 +105,7 @@ def test_approximates_ground_truth(self, t: Optional[float], order: int): @pytest.mark.parametrize(("jit", "normalize"), [(False, True), (True, False)]) def test_directed_graph(self, jit: bool, normalize: bool): - def create_graph(G: jnp.ndarray) -> graph.Graph: + def create_graph(G: jax.Array) -> graph.Graph: return geodesic.Geodesic.from_graph(G, directed=True, normalize=normalize) G = _graphs.random_graph(16, p=0.25, directed=True) @@ -127,7 +126,7 @@ def create_graph(G: jnp.ndarray) -> graph.Graph: @pytest.mark.parametrize("normalize", [False, True]) def test_normalize_laplacian(self, directed: bool, normalize: bool): - def laplacian(G: jnp.ndarray) -> jnp.ndarray: + def laplacian(G: jax.Array) -> jax.Array: if directed: G = G + G.T diff --git a/tests/geometry/graph_test.py b/tests/geometry/graph_test.py index ab638a592..f07db7fca 100644 --- a/tests/geometry/graph_test.py +++ b/tests/geometry/graph_test.py @@ -11,7 +11,7 @@ # 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 typing import Literal, Optional, Tuple +from typing import Literal import networkx as nx from networkx.generators import balanced_tree @@ -97,7 +97,7 @@ def test_approximates_ground_truth( t=[1e-4, 1e-5], only_fast=0, ) - def test_crank_nicolson_more_stable(self, t: Optional[float], n_steps: int): + def test_crank_nicolson_more_stable(self, t: float | None, n_steps: int): tol = 5 * t G = nx.linalg.adjacency_matrix(balanced_tree(r=2, h=5)) G = jnp.asarray(G.toarray()) @@ -121,7 +121,7 @@ def test_crank_nicolson_more_stable(self, t: Optional[float], n_steps: int): @pytest.mark.parametrize(("jit", "normalize"), [(False, True), (True, False)]) def test_directed_graph(self, jit: bool, normalize: bool): - def create_graph(G: jnp.ndarray) -> graph.Graph: + def create_graph(G: jax.Array) -> graph.Graph: return graph.Graph.from_graph(G, directed=True, normalize=normalize) G = _graphs.random_graph(16, p=0.25, directed=True) @@ -142,7 +142,7 @@ def create_graph(G: jnp.ndarray) -> graph.Graph: @pytest.mark.parametrize("normalize", [False, True]) def test_normalize_laplacian(self, directed: bool, normalize: bool): - def laplacian(G: jnp.ndarray) -> jnp.ndarray: + def laplacian(G: jax.Array) -> jax.Array: if directed: G = G + G.T @@ -211,8 +211,8 @@ def test_dense_graph_differentiability( ): def callback( - data: jnp.ndarray, rows: jnp.ndarray, cols: jnp.ndarray, - shape: Tuple[int, int] + data: jax.Array, rows: jax.Array, cols: jax.Array, shape: tuple[int, + int] ) -> float: G = jesp.BCOO((data, jnp.c_[rows, cols]), shape=shape).todense() diff --git a/tests/geometry/lr_cost_test.py b/tests/geometry/lr_cost_test.py index 5667dd9d6..3ad70f39c 100644 --- a/tests/geometry/lr_cost_test.py +++ b/tests/geometry/lr_cost_test.py @@ -11,7 +11,7 @@ # 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 typing import Callable, Optional, Tuple, Union +from collections.abc import Callable import pytest @@ -50,7 +50,7 @@ def test_apply(self, rng: jax.Array): def test_conversion_pointcloud( self, rng: jax.Array, - scale_cost: Union[str, float], + scale_cost: str | float, ): """Test conversion from PointCloud to LRCGeometry.""" cost_fn = costs.SqEuclidean() @@ -127,8 +127,8 @@ def test_apply_squared(self, rng: jax.Array): @pytest.mark.parametrize("bias", [(0, 0), (4, 5)]) @pytest.mark.parametrize("scale_factor", [(1, 1), (2, 3)]) def test_add_lr_geoms( - self, rng: jax.Array, bias: Tuple[float, float], - scale_factor: Tuple[float, float] + self, rng: jax.Array, bias: tuple[float, float], + scale_factor: tuple[float, float] ): """Test application of cost to vec or matrix.""" n, m, r, q = 17, 11, 7, 2 @@ -167,8 +167,7 @@ def test_add_lr_geoms( @pytest.mark.parametrize(("scale", "scale_cost", "epsilon"), [(0.1, "mean", None), (0.9, "max_cost", 1e-2)]) def test_add_lr_geoms_scale_factor( - self, rng: jax.Array, scale: float, scale_cost: str, - epsilon: Optional[float] + self, rng: jax.Array, scale: float, scale_cost: str, epsilon: float | None ): n, d = 71, 2 rng1, rng2 = jr.split(rng, 2) @@ -190,7 +189,7 @@ def test_add_lr_geoms_scale_factor( @pytest.mark.parametrize("axis", [0, 1]) @pytest.mark.parametrize("fn", [lambda x: x + 10, lambda x: x * 2]) def test_apply_affine_function_efficient( - self, rng: jax.Array, fn: Callable[[jnp.ndarray], jnp.ndarray], axis: int + self, rng: jax.Array, fn: Callable[[jax.Array], jax.Array], axis: int ): n, m, d = 21, 13, 3 rngs = jr.split(rng, 3) @@ -271,7 +270,7 @@ def test_geometry_to_lr(self, rng: jax.Array, rank: int, tol: float): only_fast=1 ) def test_point_cloud_to_lr( - self, rng: jax.Array, batch_size: Optional[int], scale_cost: Optional[str] + self, rng: jax.Array, batch_size: int | None, scale_cost: str | None ): rank, tol = 7, 1e-1 rng1, rng2 = jr.split(rng, 2) diff --git a/tests/geometry/lr_kernel_test.py b/tests/geometry/lr_kernel_test.py index 33dc140e9..3969fe7e5 100644 --- a/tests/geometry/lr_kernel_test.py +++ b/tests/geometry/lr_kernel_test.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional +from typing import Literal import pytest @@ -93,7 +93,7 @@ def test_sinkhorn_approximation( rng: jax.Array, kernel: Literal["gaussian", "arccos"], std: float, - n: Optional[int], + n: int | None, ): rng, rng1, rng2 = jr.split(rng, 3) x = jr.normal(rng1, (83, 5)) diff --git a/tests/geometry/pointcloud_test.py b/tests/geometry/pointcloud_test.py index 82ff071ee..73143bf92 100644 --- a/tests/geometry/pointcloud_test.py +++ b/tests/geometry/pointcloud_test.py @@ -11,7 +11,6 @@ # 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 typing import Union import pytest @@ -25,7 +24,7 @@ class NonSymCost(costs.CostFn): - def __call__(self, x: jnp.ndarray, y: jnp.ndarray) -> float: + def __call__(self, x: jax.Array, y: jax.Array) -> float: z = x - y return jnp.sum(z ** 2 * (jnp.sign(z) + 0.5) ** 2) @@ -128,7 +127,7 @@ class TestPointCloudCosineConversion: @pytest.mark.parametrize("scale_cost", ["mean", "median", "max_cost", 41]) def test_cosine_to_sqeucl_conversion( - self, rng: jax.Array, scale_cost: Union[str, float] + self, rng: jax.Array, scale_cost: str | float ): rng1, rng2 = jr.split(rng, 2) x = jr.normal(rng1, shape=(101, 4)) @@ -159,7 +158,7 @@ def test_cosine_to_sqeucl_conversion( @pytest.mark.parametrize("scale_cost", ["mean", "median", "max_cost", 2.0]) @pytest.mark.parametrize("axis", [0, 1]) def test_apply_cost_cosine_to_sqeucl( - self, rng: jax.Array, axis: int, scale_cost: Union[str, float] + self, rng: jax.Array, axis: int, scale_cost: str | float ): rng1, rng2 = jr.split(rng, 2) x = jr.normal(rng1, shape=(17, 5)) diff --git a/tests/geometry/regularizers_test.py b/tests/geometry/regularizers_test.py index c06f64607..5acf948c2 100644 --- a/tests/geometry/regularizers_test.py +++ b/tests/geometry/regularizers_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional import lineax as lx @@ -37,7 +36,7 @@ def test_moreau_envelope( rng: jax.Array, tau: float, reg: regularizers.ProximalOperator, - lam: Optional[float], + lam: float | None, ): tol = 1e-5 x = jr.normal(rng, (32, self.D)) @@ -115,7 +114,7 @@ def test_quad_properties( is_factor: bool, ): - def loss(reg: regularizers.ProximalOperator, x: jnp.ndarray) -> float: + def loss(reg: regularizers.ProximalOperator, x: jax.Array) -> float: return jnp.mean(jax.vmap(reg)(x)) def test_properties(reg: regularizers.ProximalOperator) -> None: diff --git a/tests/geometry/scaling_cost_test.py b/tests/geometry/scaling_cost_test.py index 29d4fae48..1f365a921 100644 --- a/tests/geometry/scaling_cost_test.py +++ b/tests/geometry/scaling_cost_test.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Optional, Union import pytest @@ -31,16 +30,16 @@ @dataclasses.dataclass(frozen=True) class ScaleCostData: """Inputs shared by every scale-cost test.""" - x: jnp.ndarray # (n, dim) - y: jnp.ndarray # (m, dim) - a: jnp.ndarray # (n,), deliberately not normalized - b: jnp.ndarray # (m,), deliberately not normalized - vec: jnp.ndarray # (m,) - cost1: jnp.ndarray # (n, 2), low-rank cost factor - cost2: jnp.ndarray # (m, 2), low-rank cost factor + x: jax.Array # (n, dim) + y: jax.Array # (m, dim) + a: jax.Array # (n,), deliberately not normalized + b: jax.Array # (m,), deliberately not normalized + vec: jax.Array # (m,) + cost1: jax.Array # (n, 2), low-rank cost factor + cost2: jax.Array # (m, 2), low-rank cost factor @property - def cost(self) -> jnp.ndarray: + def cost(self) -> jax.Array: """Squared Euclidean cost matrix between :attr:`x` and :attr:`y`.""" return ((self.x[:, None, :] - self.y[None, :, :]) ** 2).sum(-1) @@ -69,14 +68,13 @@ class TestScaleCost: only_fast=[0, -3], ) def test_scale_cost_pointcloud( - self, data: ScaleCostData, scale: Union[str, float], - batch_size: Optional[int] + self, data: ScaleCostData, scale: str | float, batch_size: int | None ): """Test various scale cost options for pointcloud.""" def apply_sinkhorn( - x: jnp.ndarray, y: jnp.ndarray, a: jnp.ndarray, b: jnp.ndarray, - scale_cost: Union[str, float] + x: jax.Array, y: jax.Array, a: jax.Array, b: jax.Array, + scale_cost: str | float ): geom = pointcloud.PointCloud(x, y, epsilon=EPS, scale_cost=scale_cost) prob = linear_problem.LinearProblem(geom, a, b) @@ -112,7 +110,7 @@ def apply_sinkhorn( "scale", ["mean", "max_cost", "max_norm", "max_bound", 100.0] ) def test_online_matches_offline_pointcloud( - self, data: ScaleCostData, scale: Union[str, float] + self, data: ScaleCostData, scale: str | float ): """Tests that the scale factors for online matches the ones without.""" geom0 = pointcloud.PointCloud( @@ -138,14 +136,11 @@ def test_online_matches_offline_pointcloud( @pytest.mark.fast.with_args( "scale", ["median", "mean", "max_cost", 100.0], only_fast=1 ) - def test_scale_cost_geometry( - self, data: ScaleCostData, scale: Union[str, float] - ): + def test_scale_cost_geometry(self, data: ScaleCostData, scale: str | float): """Test various scale cost options for geometry.""" def apply_sinkhorn( - cost: jnp.ndarray, a: jnp.ndarray, b: jnp.ndarray, - scale_cost: Union[str, float] + cost: jax.Array, a: jax.Array, b: jax.Array, scale_cost: str | float ): geom = geometry.Geometry(cost, epsilon=EPS, scale_cost=scale_cost) prob = linear_problem.LinearProblem(geom, a, b) @@ -177,9 +172,7 @@ def apply_sinkhorn( @pytest.mark.fast.with_args( "scale", ["mean", "max_bound", "max_cost", 100.0], only_fast=2 ) - def test_scale_cost_low_rank( - self, data: ScaleCostData, scale: Union[str, float] - ): + def test_scale_cost_low_rank(self, data: ScaleCostData, scale: str | float): """Test various scale cost options for low rank.""" def apply_sinkhorn(cost1, cost2, scale_cost): diff --git a/tests/geometry/semidiscrete_pointcloud_test.py b/tests/geometry/semidiscrete_pointcloud_test.py index 49447b17b..241c1c8dd 100644 --- a/tests/geometry/semidiscrete_pointcloud_test.py +++ b/tests/geometry/semidiscrete_pointcloud_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional, Tuple import pytest @@ -46,7 +45,7 @@ def test_sample(self, rng: jax.Array, num_samples: int): assert pc.shape == (num_samples, m) @pytest.mark.parametrize("epsilon", [0.0, None, 1e-2]) - def epsilon(self, rng: jax.Array, epsilon: Optional[float]): + def epsilon(self, rng: jax.Array, epsilon: float | None): n, d = 32, 5 y = jr.normal(rng, (n, d)) @@ -79,9 +78,7 @@ def test_shape(self, rng: jax.Array, n: int, m: int, d: int): @pytest.mark.parametrize(("epsilon", "dtype"), [(0.0, jnp.float16), (None, jnp.bfloat16), (0.2, jnp.float32)]) - def test_dtype( - self, rng: jax.Array, epsilon: Optional[float], dtype: jnp.dtype - ): + def test_dtype(self, rng: jax.Array, epsilon: float | None, dtype: jnp.dtype): rng_data, rng_sample = jr.split(rng, 2) m, d = 15, 1 y = jr.normal(rng_data, (m, d), dtype=dtype) @@ -99,7 +96,7 @@ def test_jit(self, rng: jax.Array): @jax.jit def sample( geom: sdpc.SemidiscretePointCloud - ) -> Tuple[sdpc.SemidiscretePointCloud, pointcloud.PointCloud, jax.Array]: + ) -> tuple[sdpc.SemidiscretePointCloud, pointcloud.PointCloud, jax.Array]: pc = geom.sample(rng_sample, 32) return geom, pc, geom.epsilon diff --git a/tests/initializers/_problems.py b/tests/initializers/_problems.py index c8e43ccae..ebd2f0c74 100644 --- a/tests/initializers/_problems.py +++ b/tests/initializers/_problems.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. """Problems shared by the initializer tests.""" -from typing import Optional import jax import jax.numpy as jnp @@ -29,7 +28,7 @@ def create_ot_problem( n: int, m: int, epsilon: float = 1e-2, - batch_size: Optional[int] = None, + batch_size: int | None = None, ) -> linear_problem.LinearProblem: """Two well-separated Gaussian clouds in 2D, with uniform marginals.""" rng_x, rng_y = jr.split(rng) diff --git a/tests/initializers/linear/sinkhorn_init_test.py b/tests/initializers/linear/sinkhorn_init_test.py index b4b693a8a..296ee297b 100644 --- a/tests/initializers/linear/sinkhorn_init_test.py +++ b/tests/initializers/linear/sinkhorn_init_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional import pytest @@ -31,7 +30,7 @@ def create_sorting_problem( rng: jax.Array, n: int, epsilon: float = 1e-2, - batch_size: Optional[int] = None + batch_size: int | None = None ) -> linear_problem.LinearProblem: # define ot problem x_init = jnp.array([-1.0, 0.0, 0.22]) diff --git a/tests/initializers/neural/meta_initializer_test.py b/tests/initializers/neural/meta_initializer_test.py index 4779c012b..02c2f41f5 100644 --- a/tests/initializers/neural/meta_initializer_test.py +++ b/tests/initializers/neural/meta_initializer_test.py @@ -15,7 +15,6 @@ import pytest import jax -import jax.numpy as jnp from flax import linen as nn @@ -30,7 +29,7 @@ class MetaMLP(nn.Module): num_hidden_layers: int = 3 @nn.compact - def __call__(self, z: jnp.ndarray) -> jnp.ndarray: + def __call__(self, z: jax.Array) -> jax.Array: for _ in range(self.num_hidden_layers): z = nn.relu(nn.Dense(self.num_hidden_units)(z)) return nn.Dense(self.potential_size)(z) diff --git a/tests/math/matrix_square_root_test.py b/tests/math/matrix_square_root_test.py index a241e40ff..07ba86b42 100644 --- a/tests/math/matrix_square_root_test.py +++ b/tests/math/matrix_square_root_test.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import pytest @@ -40,9 +41,9 @@ def _get_random_spd_matrix(dim: int, rng: jax.Array): def _get_test_fn( - fn: Callable[[jnp.ndarray], jnp.ndarray], dim: int, rng: jax.Array, + fn: Callable[[jax.Array], jax.Array], dim: int, rng: jax.Array, **kwargs: Any -) -> Callable[[jnp.ndarray], jnp.ndarray]: +) -> Callable[[jax.Array], jax.Array]: # We want to test gradients of a function fn that maps positive definite # matrices to positive definite matrices by comparing them to finite # difference approximations. We'll do so via a test function that @@ -57,7 +58,7 @@ def _get_test_fn( unit = jr.normal(rng3, shape=(dim, dim)) unit /= jnp.sqrt(jnp.sum(unit ** 2)) - def _test_fn(x: jnp.ndarray, **kwargs: Any) -> jnp.ndarray: + def _test_fn(x: jax.Array, **kwargs: Any) -> jax.Array: # m is the product of 2 symmetric, positive definite matrices # so it will be positive definite but not necessarily symmetric m = jnp.matmul(m0, m1 + x * dx) @@ -66,7 +67,7 @@ def _test_fn(x: jnp.ndarray, **kwargs: Any) -> jnp.ndarray: return _test_fn -def _sqrt_plus_inv_sqrt(x: jnp.ndarray) -> jnp.ndarray: +def _sqrt_plus_inv_sqrt(x: jax.Array) -> jax.Array: sqrtm = matrix_square_root.sqrtm(x) return sqrtm[0] + sqrtm[1] @@ -77,10 +78,10 @@ def _sqrt_plus_inv_sqrt(x: jnp.ndarray) -> jnp.ndarray: @dataclasses.dataclass(frozen=True) class Sylvester: """A Sylvester system ``a @ x - x @ b == c``, with a known solution.""" - a: jnp.ndarray # (2, m, m) - b: jnp.ndarray # (2, n, n) - x: jnp.ndarray # (2, m, n), the solution - c: jnp.ndarray # (2, m, n) + a: jax.Array # (2, m, m) + b: jax.Array # (2, n, n) + x: jax.Array # (2, m, n), the solution + c: jax.Array # (2, m, n) @pytest.fixture(scope="module") diff --git a/tests/neural/data/ot_dataloader_test.py b/tests/neural/data/ot_dataloader_test.py index 998ba82ba..1dfcd8d14 100644 --- a/tests/neural/data/ot_dataloader_test.py +++ b/tests/neural/data/ot_dataloader_test.py @@ -11,7 +11,7 @@ # 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 typing import Iterable, Tuple +from collections.abc import Iterable import jax import jax.random as jr @@ -22,8 +22,8 @@ def _get_dataset( - rng: jax.Array, shape: Tuple[int, ...] -) -> Iterable[Tuple[jax.Array, jax.Array]]: + rng: jax.Array, shape: tuple[int, ...] +) -> Iterable[tuple[jax.Array, jax.Array]]: while True: rng, rng_x0, rng_x1 = jr.split(rng, 3) x0 = jr.normal(rng_x0, shape) diff --git a/tests/neural/data/semidiscrete_dataloader_test.py b/tests/neural/data/semidiscrete_dataloader_test.py index f33df98d8..60b6fd426 100644 --- a/tests/neural/data/semidiscrete_dataloader_test.py +++ b/tests/neural/data/semidiscrete_dataloader_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional, Tuple import pytest @@ -29,7 +28,7 @@ def _solve_semidiscrete( - rng: jax.Array, *, shape: Tuple[int, ...], epsilon: float + rng: jax.Array, *, shape: tuple[int, ...], epsilon: float ) -> semidiscrete.SemidiscreteOutput: rng_data, rng_solve = jr.split(rng, 2) y = jr.normal(rng_data, shape) @@ -116,7 +115,7 @@ def test_subset_size_threshold( assert tgt.shape == (batch_size, d) @pytest.mark.parametrize("epsilon", [0.0, 1e-2, None]) - def test_sharding(self, rng: jax.Array, epsilon: Optional[float]): + def test_sharding(self, rng: jax.Array, epsilon: float | None): m, d = 11, 4 batch_size = 11 rng_solve, rng_dl = jr.split(rng, 2) diff --git a/tests/neural/methods/flow_matching_test.py b/tests/neural/methods/flow_matching_test.py index 9159e9ce2..64d7ef9ec 100644 --- a/tests/neural/methods/flow_matching_test.py +++ b/tests/neural/methods/flow_matching_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Dict, Literal, Optional, Tuple, Union +from typing import Literal import pytest @@ -32,10 +32,10 @@ def _prepare_batch( rng: jax.Array, *, - shape: Tuple[int, ...], - num_classes: Optional[int] = None, + shape: tuple[int, ...], + num_classes: int | None = None, x0_stddev: float = 1.0, -) -> Dict[Literal["t", "x_t", "v_t", "cond", "x0", "x1"], jax.Array]: +) -> dict[Literal["t", "x_t", "v_t", "cond", "x0", "x1"], jax.Array]: rng_t, rng_x0, rng_x1, rng_cond = jr.split(rng, 4) batch_size, *_ = shape x0 = jr.normal(rng_x0, shape) * x0_stddev @@ -92,7 +92,7 @@ def test_ema_callback(self, rng: jax.Array): @pytest.mark.parametrize(("num_steps", "reverse"), [(None, False), (4, True)]) def test_evaluate_vf( - self, rng: jax.Array, num_steps: Optional[int], reverse: bool + self, rng: jax.Array, num_steps: int | None, reverse: bool ): batch_size, dim = 2, 3 batch = _prepare_batch(rng, shape=(batch_size, dim)) @@ -114,9 +114,7 @@ def test_evaluate_vf( assert sol.ys.shape == (batch_size, 1, dim) @pytest.mark.parametrize("num_steps", [None, 3]) - def test_evaluate_vf_save_extra( - self, rng: jax.Array, num_steps: Optional[int] - ): + def test_evaluate_vf_save_extra(self, rng: jax.Array, num_steps: int | None): batch_size, dim = 5, 3 ode_max_steps, vel_save_steps = 16, 4 batch = _prepare_batch(rng, shape=(batch_size, dim)) @@ -145,8 +143,7 @@ def test_evaluate_vf_save_extra( @pytest.mark.parametrize("drop_last_velocity", [None, False, True]) @pytest.mark.parametrize("ts", [3, tuple(jnp.linspace(0.0, 1.0, 5).tolist())]) def test_curvature( - self, rng: jax.Array, ts: Union[int, jax.Array], - drop_last_velocity: Optional[bool] + self, rng: jax.Array, ts: int | jax.Array, drop_last_velocity: bool | None ): dim = 4 batch = _prepare_batch(rng, shape=(1, dim)) diff --git a/tests/neural/methods/monge_gap_test.py b/tests/neural/methods/monge_gap_test.py index 002a09e01..bd63a5443 100644 --- a/tests/neural/methods/monge_gap_test.py +++ b/tests/neural/methods/monge_gap_test.py @@ -11,12 +11,10 @@ # 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 typing import Optional import pytest import jax -import jax.numpy as jnp import jax.random as jr import numpy as np @@ -141,9 +139,9 @@ def test_map_estimator_convergence(self): # define the fitting loss and the regularizer def fitting_loss( - samples: jnp.ndarray, - mapped_samples: jnp.ndarray, - ) -> Optional[float]: + samples: jax.Array, + mapped_samples: jax.Array, + ) -> float | None: r"""Sinkhorn divergence fitting loss.""" div, _ = sinkhorn_divergence.sinkdiv( x=samples, diff --git a/tests/neural/methods/neuraldual_test.py b/tests/neural/methods/neuraldual_test.py index 65b78f7dd..9525a469d 100644 --- a/tests/neural/methods/neuraldual_test.py +++ b/tests/neural/methods/neuraldual_test.py @@ -11,7 +11,7 @@ # 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 typing import Optional, Sequence, Tuple +from collections.abc import Sequence import pytest @@ -25,12 +25,12 @@ from ott.neural.networks import icnn, potentials from ott.neural.networks.layers import conjugate -ModelPair_t = Tuple[nnx.Module, nnx.Module] -DatasetPair_t = Tuple[datasets.Dataset, datasets.Dataset] +ModelPair_t = tuple[nnx.Module, nnx.Module] +DatasetPair_t = tuple[datasets.Dataset, datasets.Dataset] @pytest.fixture(params=[("simple", "circle")]) -def ds(rng: jax.Array, request: Tuple[str, str]) -> DatasetPair_t: +def ds(rng: jax.Array, request: tuple[str, str]) -> DatasetPair_t: train_dataset, valid_dataset, _ = datasets.create_gaussian_mixture_samplers( request.param[0], request.param[1], @@ -76,7 +76,7 @@ def test_neural_dual_convergence( neural_models: ModelPair_t, back_and_forth: bool, amortization_loss: str, - conjugate_solver: Optional[conjugate.FenchelConjugateSolver], + conjugate_solver: conjugate.FenchelConjugateSolver | None, ): """Tests convergence of learning the Kantorovich dual using ICNNs.""" diff --git a/tests/neural/networks/velocity_field/unet_test.py b/tests/neural/networks/velocity_field/unet_test.py index d4df071ba..d043d60ac 100644 --- a/tests/neural/networks/velocity_field/unet_test.py +++ b/tests/neural/networks/velocity_field/unet_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional, Tuple import pytest @@ -27,9 +26,9 @@ def _prepare_inputs( rng: jax.Array, *, - shape: Tuple[int, ...], - num_classes: Optional[int] = None, -) -> Tuple[jax.Array, jax.Array, Optional[jax.Array]]: + shape: tuple[int, ...], + num_classes: int | None = None, +) -> tuple[jax.Array, jax.Array, jax.Array | None]: rng_t, rng_x, rng_cond = jr.split(rng, 3) batch_size, *_ = shape t = jr.uniform(rng_t, (batch_size,)) diff --git a/tests/problems/linear/potentials_test.py b/tests/problems/linear/potentials_test.py index 7b8c28a29..a1023abba 100644 --- a/tests/problems/linear/potentials_test.py +++ b/tests/problems/linear/potentials_test.py @@ -242,11 +242,11 @@ def test_potentials_diff_param_costs( self, rng: jax.Array, reg: regularizers.ProximalOperator ): - def proj(matrix: jnp.ndarray) -> jnp.ndarray: + def proj(matrix: jax.Array) -> jax.Array: u, _, v_h = jnp.linalg.svd(matrix, full_matrices=False) return u.dot(v_h) - def create_cost(A: jnp.ndarray) -> costs.RegTICost: + def create_cost(A: jax.Array) -> costs.RegTICost: A = lx.MatrixLinearOperator(A) orth = regularizers.Orthogonal(reg, A=A) return costs.RegTICost(orth, lam=1.0) diff --git a/tests/solvers/linear/continuous_barycenter_test.py b/tests/solvers/linear/continuous_barycenter_test.py index d0a00bc1c..6966213ca 100644 --- a/tests/solvers/linear/continuous_barycenter_test.py +++ b/tests/solvers/linear/continuous_barycenter_test.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Tuple import pytest @@ -30,7 +29,7 @@ means_and_covs_to_x = jax.vmap(costs.mean_and_cov_to_x, in_axes=[0, 0, None]) -def is_positive_semidefinite(c: jnp.ndarray) -> bool: +def is_positive_semidefinite(c: jax.Array) -> bool: # GPU friendly, eigvals not implemented for non-symmetric matrices w = jnp.linalg.eigvalsh((c + c.T) / 2.0) return jnp.all(w >= 0) @@ -129,10 +128,10 @@ def test_barycenter_jit(self, rng: jax.Array, segment_before: bool): @functools.partial(jax.jit, static_argnums=(2, 3)) def barycenter( - y: jnp.ndarray, - b: jnp.ndarray, + y: jax.Array, + b: jax.Array, segment_before: bool, - num_per_segment: Tuple[int, ...], + num_per_segment: tuple[int, ...], ) -> cb.FreeBarycenterState: if segment_before: y, b, num_per_segment = segment.segment_point_cloud( diff --git a/tests/solvers/linear/semidiscrete_test.py b/tests/solvers/linear/semidiscrete_test.py index 60535ac1b..73b37a5cd 100644 --- a/tests/solvers/linear/semidiscrete_test.py +++ b/tests/solvers/linear/semidiscrete_test.py @@ -11,7 +11,7 @@ # 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 typing import Any, Optional +from typing import Any import pytest @@ -36,7 +36,7 @@ def _random_problem( *, m: int, d: int, - dtype: Optional[jnp.dtype] = None, + dtype: jnp.dtype | None = None, **kwargs: Any ) -> sdlp.SemidiscreteLinearProblem: rng_b, rng_y = jr.split(rng, 2) @@ -54,7 +54,7 @@ class TestSemidiscreteSolver: @pytest.mark.parametrize("n", [20, 31]) @pytest.mark.parametrize("epsilon", [0.0, 1e-3, 1e-2, 1e-1, None]) def test_custom_gradient_semidiscrete_loss( - self, rng: jax.Array, n: int, epsilon: Optional[float] + self, rng: jax.Array, n: int, epsilon: float | None ): def semidiscrete_loss( @@ -85,9 +85,7 @@ def semidiscrete_loss( @pytest.mark.parametrize(("dtype", "epsilon"), [(jnp.float16, 0.0), (jnp.bfloat16, 0.5), (jnp.float32, None)]) - def test_dtype( - self, rng: jax.Array, dtype: jnp.dtype, epsilon: Optional[float] - ): + def test_dtype(self, rng: jax.Array, dtype: jnp.dtype, epsilon: float | None): m, d = 22, 3 rng_prob, rng_solver, rng_sample = jr.split(rng, 3) prob = _random_problem(rng_prob, m=m, d=d, epsilon=epsilon, dtype=dtype) @@ -135,7 +133,7 @@ def print_state(state: semidiscrete.SemidiscreteState) -> None: assert actual.err == "" @pytest.mark.parametrize("epsilon", [0.0, 1e-2, None]) - def test_epsilon(self, rng: jax.Array, epsilon: Optional[float]): + def test_epsilon(self, rng: jax.Array, epsilon: float | None): rng_prob, rng_solver, rng_sample = jr.split(rng, 3) prob = _random_problem(rng_prob, m=15, d=4, epsilon=epsilon) @@ -161,7 +159,7 @@ def test_epsilon(self, rng: jax.Array, epsilon: Optional[float]): assert isinstance(out_sampled, semidiscrete.HardAssignmentOutput) @pytest.mark.parametrize("epsilon", [1e-1, None]) - def test_match_with_finiteOT(self, rng: jax.Array, epsilon: Optional[float]): + def test_match_with_finiteOT(self, rng: jax.Array, epsilon: float | None): rng_solver, rng_sample, rng_b, rng_y = jr.split(rng, 4) m, d = 8, 2 b = jr.uniform(rng_b, (m,)) + 1. # balanced distribution helps converge @@ -200,7 +198,7 @@ def test_match_with_finiteOT(self, rng: jax.Array, epsilon: Optional[float]): np.testing.assert_allclose(g_ot, g_sd, rtol=1e-2, atol=1e-1) @pytest.mark.parametrize("epsilon", [0.0, 1e-2, None]) - def test_initial_potential(self, rng: jax.Array, epsilon: Optional[float]): + def test_initial_potential(self, rng: jax.Array, epsilon: float | None): rng_prob, rng_solver = jr.split(rng, 2) prob = _random_problem(rng_prob, m=32, d=3, epsilon=epsilon) @@ -232,7 +230,7 @@ def test_solver_wrapper(self, rng: jax.Array): @pytest.mark.parametrize(("n", "epsilon"), [(17, 0.0), (20, 1e-3), (35, None)]) - def test_output(self, rng: jax.Array, n: int, epsilon: Optional[float]): + def test_output(self, rng: jax.Array, n: int, epsilon: float | None): m, d = 32, 3 rng_prob, rng_solver, rng_sample = jr.split(rng, 3) prob = _random_problem(rng_prob, m=m, d=d, epsilon=epsilon) @@ -280,7 +278,7 @@ def test_output(self, rng: jax.Array, n: int, epsilon: Optional[float]): @pytest.mark.parametrize("epsilon_dp", [None, 0.1]) @pytest.mark.parametrize("epsilon_prob", [0.0, 1e-2]) def test_sd_dual_potentials( - self, rng: jax.Array, epsilon_prob: float, epsilon_dp: Optional[float] + self, rng: jax.Array, epsilon_prob: float, epsilon_dp: float | None ): rng_prob, rng_solver, rng_sample = jr.split(rng, 3) diff --git a/tests/solvers/linear/sinkhorn_diff_test.py b/tests/solvers/linear/sinkhorn_diff_test.py index 94238706c..be1bd25b0 100644 --- a/tests/solvers/linear/sinkhorn_diff_test.py +++ b/tests/solvers/linear/sinkhorn_diff_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Callable, List, Optional, Tuple +from collections.abc import Callable import pytest @@ -45,7 +45,7 @@ def test_implicit_differentiation_versus_autodiff( ): epsilon = 5e-2 - def loss_g(a: jnp.ndarray, x: jnp.ndarray, implicit: bool = True) -> float: + def loss_g(a: jax.Array, x: jax.Array, implicit: bool = True) -> float: implicit_diff = implicit_lib.ImplicitDiff() if implicit else None geom = geometry.Geometry( cost_matrix=jnp.sum(x ** 2, axis=1)[:, jnp.newaxis] + @@ -61,9 +61,7 @@ def loss_g(a: jnp.ndarray, x: jnp.ndarray, implicit: bool = True) -> float: ) return solver(prob).reg_ot_cost - def loss_pcg( - a: jnp.ndarray, x: jnp.ndarray, implicit: bool = True - ) -> float: + def loss_pcg(a: jax.Array, x: jax.Array, implicit: bool = True) -> float: implicit_diff = implicit_lib.ImplicitDiff() if implicit else None geom = pointcloud.PointCloud(x, clouds.y, epsilon=epsilon) prob = linear_problem.LinearProblem( @@ -136,7 +134,7 @@ class TestSinkhornJacobian: only_fast=0, ) def test_autograd_sinkhorn( - self, rng: jax.Array, lse_mode: bool, shape_data: Tuple[int, int] + self, rng: jax.Array, lse_mode: bool, shape_data: tuple[int, int] ): """Test gradient w.r.t. probability weights.""" n, m = shape_data @@ -153,7 +151,7 @@ def test_autograd_sinkhorn( a = a / jnp.sum(a) b = b / jnp.sum(b) - def reg_ot(a: jnp.ndarray, b: jnp.ndarray) -> float: + def reg_ot(a: jax.Array, b: jax.Array) -> float: geom = pointcloud.PointCloud(x, y, epsilon=1e-1) prob = linear_problem.LinearProblem(geom, a=a, b=b) solver = sinkhorn.Sinkhorn(lse_mode=lse_mode) @@ -180,7 +178,7 @@ def reg_ot(a: jnp.ndarray, b: jnp.ndarray) -> float: @pytest.mark.parametrize(("lse_mode", "shape_data"), [(True, (7, 9)), (False, (11, 5))]) def test_gradient_sinkhorn_geometry( - self, rng: jax.Array, lse_mode: bool, shape_data: Tuple[int, int] + self, rng: jax.Array, lse_mode: bool, shape_data: tuple[int, int] ): """Test gradient w.r.t. cost matrix.""" n, m = shape_data @@ -190,7 +188,7 @@ def test_gradient_sinkhorn_geometry( delta = delta / jnp.sqrt(jnp.vdot(delta, delta)) eps = 1e-3 # perturbation magnitude - def loss_fn(cm: jnp.ndarray): + def loss_fn(cm: jax.Array): geom = geometry.Geometry(cm, epsilon=0.5) prob = linear_problem.LinearProblem(geom) solver = sinkhorn.Sinkhorn(lse_mode=lse_mode) @@ -262,8 +260,8 @@ def test_gradient_sinkhorn_euclidean( # Adding some near-zero distances to test proper handling with p_norm=1. y = y.at[0].set(x[0, :] + 1e-3) - def loss_fn(x: jnp.ndarray, - y: jnp.ndarray) -> Tuple[float, sinkhorn.SinkhornOutput]: + def loss_fn(x: jax.Array, + y: jax.Array) -> tuple[float, sinkhorn.SinkhornOutput]: implicit_diff = implicit_lib.ImplicitDiff() if implicit else None geom = pointcloud.PointCloud(x, y, epsilon=epsilon, cost_fn=cost_fn) prob = linear_problem.LinearProblem(geom, a, b) @@ -318,7 +316,7 @@ def loss_fn(x: jnp.ndarray, def test_autoepsilon_differentiability(self, rng: jax.Array): cost = jr.uniform(rng, (15, 17)) - def reg_ot_cost(c: jnp.ndarray) -> float: + def reg_ot_cost(c: jax.Array) -> float: geom = geometry.Geometry(c, epsilon=None) # auto epsilon prob = linear_problem.LinearProblem(geom) return sinkhorn.Sinkhorn()(prob).reg_ot_cost @@ -329,7 +327,7 @@ def reg_ot_cost(c: jnp.ndarray) -> float: @pytest.mark.fast() def test_differentiability_with_jit(self, rng: jax.Array): - def reg_ot_cost(c: jnp.ndarray) -> float: + def reg_ot_cost(c: jax.Array) -> float: geom = geometry.Geometry(c, epsilon=1e-2) prob = linear_problem.LinearProblem(geom) return sinkhorn.Sinkhorn()(prob).reg_ot_cost @@ -381,7 +379,7 @@ def test_apply_transport_jacobian( # general rule, even more so when using backprop. epsilon = 0.01 if lse_mode else 0.1 - def apply_ot(a: jnp.ndarray, x: jnp.ndarray, implicit: bool) -> jnp.ndarray: + def apply_ot(a: jax.Array, x: jax.Array, implicit: bool) -> jax.Array: geom = pointcloud.PointCloud(x, y, epsilon=epsilon) prob = linear_problem.LinearProblem(geom, a, b, tau_a=tau_a, tau_b=tau_b) @@ -455,7 +453,7 @@ def apply_ot(a: jnp.ndarray, x: jnp.ndarray, implicit: bool) -> jnp.ndarray: ) def test_potential_jacobian_sinkhorn( self, rng: jax.Array, lse_mode: bool, tau_a: float, tau_b: float, - shape: Tuple[int, int], arg: int + shape: tuple[int, int], arg: int ): """Test Jacobian of optimal potential w.r.t. weights and locations.""" atol = 1e-2 if lse_mode else 5e-2 # lower tolerance for lse mode. @@ -482,7 +480,7 @@ def test_potential_jacobian_sinkhorn( # with small epsilon when differentiating. epsilon = 0.01 if lse_mode else 0.1 - def loss_from_potential(a: jnp.ndarray, x: jnp.ndarray, implicit: bool): + def loss_from_potential(a: jax.Array, x: jax.Array, implicit: bool): geom = pointcloud.PointCloud(x, y, epsilon=epsilon) prob = linear_problem.LinearProblem(geom, a, b, tau_a=tau_a, tau_b=tau_b) @@ -551,7 +549,7 @@ def test_diff_sinkhorn_x_grid_x_perturbation( a = a.ravel() / jnp.sum(a) b = b.ravel() / jnp.sum(b) - def reg_ot(x: List[jnp.ndarray]) -> float: + def reg_ot(x: list[jax.Array]) -> float: geom = grid.Grid(x=x, epsilon=1.0) prob = linear_problem.LinearProblem(geom, a=a, b=b) solver = sinkhorn.Sinkhorn(threshold=1e-1, lse_mode=lse_mode) @@ -601,7 +599,7 @@ def test_diff_sinkhorn_x_grid_weights_perturbation( b = b.ravel() / jnp.sum(b) geom = grid.Grid(x=x, epsilon=1) - def reg_ot(a: jnp.ndarray, b: jnp.ndarray) -> float: + def reg_ot(a: jax.Array, b: jax.Array) -> float: prob = linear_problem.LinearProblem(geom, a, b) solver = sinkhorn.Sinkhorn(threshold=1e-3, lse_mode=lse_mode) return solver(prob).reg_ot_cost @@ -635,7 +633,7 @@ class TestSinkhornJacobianPreconditioning: ) def test_potential_jacobian_sinkhorn_precond( self, rng: jax.Array, lse_mode: bool, tau_a: float, tau_b: float, - shape: Tuple[int, int], arg: int + shape: tuple[int, int], arg: int ): """Test Jacobian of optimal potential works across 2 precond_fun.""" atol = 1e-2 if lse_mode else 5e-2 # lower tolerance for lse mode. @@ -663,9 +661,9 @@ def test_potential_jacobian_sinkhorn_precond( epsilon = 0.05 if lse_mode else 0.1 def loss_from_potential( - a: jnp.ndarray, - x: jnp.ndarray, - precondition_fun: Optional[Callable[[jnp.ndarray], jnp.ndarray]] = None, + a: jax.Array, + x: jax.Array, + precondition_fun: Callable[[jax.Array], jax.Array] | None = None, symmetric: bool = False ) -> float: geom = pointcloud.PointCloud(x, y, epsilon=epsilon) @@ -757,7 +755,7 @@ def test_hessian_sinkhorn( imp_dif = implicit_lib.ImplicitDiff(solver_kwargs=solver_kwargs) - def loss(a: jnp.ndarray, x: jnp.ndarray, implicit: bool = True): + def loss(a: jax.Array, x: jax.Array, implicit: bool = True): geom = pointcloud.PointCloud(x, y, epsilon=epsilon) prob = linear_problem.LinearProblem(geom, a, b, tau_a, tau_b) implicit_diff = imp_dif if implicit else None diff --git a/tests/solvers/linear/sinkhorn_lr_test.py b/tests/solvers/linear/sinkhorn_lr_test.py index 5ab499526..746731f49 100644 --- a/tests/solvers/linear/sinkhorn_lr_test.py +++ b/tests/solvers/linear/sinkhorn_lr_test.py @@ -11,7 +11,6 @@ # 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 typing import Type import pytest @@ -46,7 +45,7 @@ class TestLRSinkhorn: ) def test_euclidean_point_cloud_lr( self, clouds: _utils.PointClouds, use_lrcgeom: bool, - initializer_class: Type[initializers_lr.LRInitializer], + initializer_class: type[initializers_lr.LRInitializer], gamma_rescale: bool, lse_mode: bool ): """Two point clouds, tested with 3 different initializations.""" diff --git a/tests/solvers/linear/sinkhorn_misc_test.py b/tests/solvers/linear/sinkhorn_misc_test.py index a5bd66be9..93e09ee03 100644 --- a/tests/solvers/linear/sinkhorn_misc_test.py +++ b/tests/solvers/linear/sinkhorn_misc_test.py @@ -11,7 +11,6 @@ # 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 typing import Optional import pytest @@ -270,7 +269,7 @@ def test_sinkhorn_unbalanced_recenter_acceleration( eps: float, tau_a: float, tau_b: float, - anderson: Optional[acceleration.AndersonAcceleration], + anderson: acceleration.AndersonAcceleration | None, ): def run_sink(*, recenter: bool) -> sinkhorn.SinkhornOutput: @@ -329,12 +328,10 @@ def assert_output_close( ) -> None: """Assert SinkhornOutputs are close.""" x = tuple( - a for a in x - if (a is not None and (isinstance(a, (jnp.ndarray, int)))) + a for a in x if (a is not None and (isinstance(a, (jax.Array, int)))) ) y = tuple( - a for a in y - if (a is not None and (isinstance(a, (jnp.ndarray, int)))) + a for a in y if (a is not None and (isinstance(a, (jax.Array, int)))) ) return chex.assert_trees_all_close(x, y, atol=1e-6, rtol=0) @@ -349,7 +346,7 @@ def test_jit_vs_non_jit_bwd( ): @jax.value_and_grad - def val_grad(a: jnp.ndarray, x: jnp.ndarray) -> float: + def val_grad(a: jax.Array, x: jax.Array) -> float: implicit_diff = implicit_lib.ImplicitDiff() if implicit else None geom = geometry.Geometry( cost_matrix=( diff --git a/tests/solvers/linear/sinkhorn_test.py b/tests/solvers/linear/sinkhorn_test.py index fbf43fdcc..14c7c2c50 100644 --- a/tests/solvers/linear/sinkhorn_test.py +++ b/tests/solvers/linear/sinkhorn_test.py @@ -14,7 +14,6 @@ import functools import io import sys -from typing import Optional import pytest @@ -509,9 +508,7 @@ def test_sinkhorn_online_memory_jit(self, rng: jax.Array): assert out.primal_cost > 0.0 @pytest.mark.fast.with_args(cost_fn=[None, costs.SqPNorm(1.6)], only_fast=0) - def test_primal_cost_grid( - self, rng: jax.Array, cost_fn: Optional[costs.CostFn] - ): + def test_primal_cost_grid(self, rng: jax.Array, cost_fn: costs.CostFn | None): """Test computation of primal / costs for Grids.""" rng_a, rng_b = jr.split(rng) ns = [6, 7, 11] @@ -679,7 +676,7 @@ def test_custom_progress_fn(self, clouds: _utils.PointClouds): ] @pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) - def test_sinkhorn_dtype(self, clouds: _utils.PointClouds, dtype: jnp.ndarray): + def test_sinkhorn_dtype(self, clouds: _utils.PointClouds, dtype: jax.Array): x = clouds.x.astype(dtype) y = clouds.y.astype(dtype) geom = pointcloud.PointCloud(x, y, epsilon=jnp.array(1e-1, dtype=dtype)) diff --git a/tests/solvers/linear/univariate_test.py b/tests/solvers/linear/univariate_test.py index b90094eb2..5a130cd42 100644 --- a/tests/solvers/linear/univariate_test.py +++ b/tests/solvers/linear/univariate_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Callable +from collections.abc import Callable import pytest @@ -82,9 +82,7 @@ def test_cdf_distance_and_sinkhorn( @jax.jit @functools.partial(jax.vmap, in_axes=[1, 1, None, None]) - def sliced_sinkhorn( - x: jnp.ndarray, y: jnp.ndarray, a: jnp.ndarray, b: jnp.ndarray - ): + def sliced_sinkhorn(x: jax.Array, y: jax.Array, a: jax.Array, b: jax.Array): geom = pointcloud.PointCloud( x[:, None], y[:, None], cost_fn=cost_fn, epsilon=1e-4 ) @@ -166,7 +164,7 @@ def test_univariate_grad( ): def univ_dist( - x: jnp.ndarray, y: jnp.ndarray, a: jnp.ndarray, b: jnp.ndarray + x: jax.Array, y: jax.Array, a: jax.Array, b: jax.Array ) -> float: geom = pointcloud.PointCloud(x[:, None], y[:, None]) prob = linear_problem.LinearProblem(geom, a=a, b=b) diff --git a/tests/solvers/quadratic/fgw_test.py b/tests/solvers/quadratic/fgw_test.py index bd1c77207..ad7583c18 100644 --- a/tests/solvers/quadratic/fgw_test.py +++ b/tests/solvers/quadratic/fgw_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import dataclasses -from typing import Literal, Tuple, Union +from typing import Literal import pytest @@ -35,9 +35,9 @@ @dataclasses.dataclass(frozen=True) class FusedClouds(_utils.QuadClouds): """:class:`~tests._utils.QuadClouds` plus the fused, inter-domain data.""" - x_2: jnp.ndarray # (n, d_xy), source of the fused term - y_2: jnp.ndarray # (m, d_xy), target of the fused term - cxy: jnp.ndarray # (n, m), inter-domain cost + x_2: jax.Array # (n, d_xy), source of the fused term + y_2: jax.Array # (m, d_xy), target of the fused term + cxy: jax.Array # (n, m), inter-domain cost @pytest.fixture(scope="module") @@ -67,7 +67,7 @@ def test_gradient_marginals_fgw_solver(self, clouds: FusedClouds, jit: bool): geom_y = pointcloud.PointCloud(clouds.y) geom_xy = pointcloud.PointCloud(clouds.x_2, clouds.y_2) - def reg_gw(a: jnp.ndarray, b: jnp.ndarray, implicit: bool): + def reg_gw(a: jax.Array, b: jax.Array, implicit: bool): prob = quadratic_problem.QuadraticProblem( geom_x, geom_y, geom_xy, fused_penalty=FUSED_PENALTY, a=a, b=b ) @@ -112,9 +112,8 @@ def test_gradient_fgw_solver_geometry( """Test gradient w.r.t. the geometries.""" def reg_gw( - x: jnp.ndarray, y: jnp.ndarray, - xy: Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]], - fused_penalty: float, a: jnp.ndarray, b: jnp.ndarray, implicit: bool + x: jax.Array, y: jax.Array, xy: jax.Array | tuple[jax.Array, jax.Array], + fused_penalty: float, a: jax.Array, b: jax.Array, implicit: bool ): if is_cost: geom_x = geometry.Geometry(cost_matrix=x) @@ -189,8 +188,8 @@ def test_gradient_fgw_solver_penalty(self, clouds: FusedClouds): lse_mode = True def reg_gw( - cx: jnp.ndarray, cy: jnp.ndarray, cxy: jnp.ndarray, - fused_penalty: float, a: jnp.ndarray, b: jnp.ndarray, implicit: bool + cx: jax.Array, cy: jax.Array, cxy: jax.Array, fused_penalty: float, + a: jax.Array, b: jax.Array, implicit: bool ) -> float: geom_x = geometry.Geometry(cost_matrix=cx) geom_y = geometry.Geometry(cost_matrix=cy) @@ -254,7 +253,7 @@ def test_fgw_lr_memory(self, rng: jax.Array, jit: bool): @pytest.mark.parametrize("cost_rank", [4, (2, 3, 4)]) def test_fgw_lr_generic_cost_matrix( - self, rng: jax.Array, cost_rank: Union[int, Tuple[int, int, int]] + self, rng: jax.Array, cost_rank: int | tuple[int, int, int] ): n, m = 20, 30 rng1, rng2, rng3, rng4 = jr.split(rng, 4) diff --git a/tests/solvers/quadratic/gw_barycenter_test.py b/tests/solvers/quadratic/gw_barycenter_test.py index 2af22df8f..5a2d9ffa0 100644 --- a/tests/solvers/quadratic/gw_barycenter_test.py +++ b/tests/solvers/quadratic/gw_barycenter_test.py @@ -11,7 +11,8 @@ # 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 typing import Any, Optional, Sequence, Tuple +from collections.abc import Sequence +from typing import Any import pytest @@ -36,7 +37,7 @@ def random_pc( n: int, d: int, rng: jax.Array, - m: Optional[int] = None, + m: int | None = None, **kwargs: Any ) -> pointcloud.PointCloud: rng1, rng2 = jr.split(rng, 2) @@ -46,9 +47,9 @@ def random_pc( @staticmethod def pad_cost_matrices( - costs: Sequence[jnp.ndarray], - shape: Optional[Tuple[int, int]] = None - ) -> Tuple[jnp.ndarray, jnp.ndarray]: + costs: Sequence[jax.Array], + shape: tuple[int, int] | None = None + ) -> tuple[jax.Array, jax.Array]: if shape is None: shape = jnp.asarray([arr.shape for arr in costs]).max() shape = (shape, shape) @@ -69,8 +70,7 @@ def pad_cost_matrices( [("sqeucl", 17, None)] # , ("kl", 22, 1e-2)] ) def test_gw_barycenter( - self, rng: jax.Array, gw_loss: str, bar_size: int, - epsilon: Optional[float] + self, rng: jax.Array, gw_loss: str, bar_size: int, epsilon: float | None ): tol = 1e-3 if gw_loss == "sqeucl" else 1e-1 num_per_segment = (13, 15, 21) @@ -139,7 +139,7 @@ def test_fgw_barycenter( ): def barycenter( - y: jnp.ndim, y_fused: jnp.ndarray, num_per_segment: Tuple[int, ...] + y: jnp.ndim, y_fused: jax.Array, num_per_segment: tuple[int, ...] ) -> gwb_solver.GWBarycenterState: bar_prob = gwb.GWBarycenterProblem( y=y, diff --git a/tests/solvers/quadratic/gw_test.py b/tests/solvers/quadratic/gw_test.py index ff101c469..6d731a4a1 100644 --- a/tests/solvers/quadratic/gw_test.py +++ b/tests/solvers/quadratic/gw_test.py @@ -11,7 +11,6 @@ # 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 typing import Tuple, Union import pytest @@ -36,7 +35,7 @@ class TestQuadraticProblem: @pytest.mark.parametrize("rank", [-1, 5, (1, 2, 3), (2, 3, 5)]) def test_quad_to_low_rank( self, clouds: _utils.QuadClouds, rng: jax.Array, as_pc: bool, - rank: Union[int, Tuple[int, ...]] + rank: int | tuple[int, ...] ): n, m, d1, d2, d = 100, 120, 4, 6, 10 rng1, rng2, rng3, rng4 = jr.split(rng, 4) @@ -158,8 +157,8 @@ def test_flag_store_errors(self, clouds: _utils.QuadClouds): def test_gradient_marginals_gw(self, clouds: _utils.QuadClouds, jit: bool): """Test gradient w.r.t. probability weights.""" - def reg_gw(a: jnp.ndarray, b: jnp.ndarray, - implicit: bool) -> Tuple[float, Tuple[jnp.ndarray, jnp.ndarray]]: + def reg_gw(a: jax.Array, b: jax.Array, + implicit: bool) -> tuple[float, tuple[jax.Array, jax.Array]]: prob = quadratic_problem.QuadraticProblem(geom_x, geom_y, a=a, b=b) implicit_diff = implicit_lib.ImplicitDiff() if implicit else None linear_solver = sinkhorn.Sinkhorn( @@ -251,8 +250,7 @@ def test_gradient_gw_geometry( """Test gradient w.r.t. the geometries.""" def reg_gw( - x: jnp.ndarray, y: jnp.ndarray, a: jnp.ndarray, b: jnp.ndarray, - implicit: bool + x: jax.Array, y: jax.Array, a: jax.Array, b: jax.Array, implicit: bool ) -> float: if is_cost: geom_x = geometry.Geometry(cost_matrix=x) @@ -414,7 +412,7 @@ def test_gw_lr_apply(self, clouds: _utils.QuadClouds, axis: int): def test_relative_epsilon( self, rng: jax.Array, - scale_cost: Union[float, str], + scale_cost: float | str, ): eps = 1e-2 rng1, rng2 = jr.split(rng, 2) @@ -523,7 +521,7 @@ def test_gwlr_unbalanced_matches_balanced( @pytest.mark.parametrize("grad", [False, True]) def test_gw_progress_fn(self, clouds: _utils.QuadClouds, grad: bool): - def callback(x: jnp.ndarray, y: jnp.ndarray): + def callback(x: jax.Array, y: jax.Array): geom_xx = pointcloud.PointCloud(x) geom_yy = pointcloud.PointCloud(y) prob = quadratic_problem.QuadraticProblem(geom_xx, geom_yy) diff --git a/tests/tools/conformal_test.py b/tests/tools/conformal_test.py index e410d127d..c54c05623 100644 --- a/tests/tools/conformal_test.py +++ b/tests/tools/conformal_test.py @@ -13,7 +13,7 @@ # limitations under the License. import functools import math -from typing import Callable, Tuple +from collections.abc import Callable import pytest @@ -30,7 +30,7 @@ def get_model_and_data( n_samples: int, target_dim: int, random_state: int = 0, -) -> Tuple[Callable[[jnp.ndarray], jnp.ndarray], Tuple[jnp.ndarray, ...]]: +) -> tuple[Callable[[jax.Array], jax.Array], tuple[jax.Array, ...]]: x, y = datasets.make_regression( n_samples=n_samples, n_features=5, @@ -57,7 +57,7 @@ def get_model_and_data( class TestOTCP: @pytest.mark.parametrize("shape", [(16, 2), (58, 9), (128, 9)]) - def test_sobol_ball_sampler(self, shape: Tuple[int, int], rng: jax.Array): + def test_sobol_ball_sampler(self, shape: tuple[int, int], rng: jax.Array): n, d = shape n_per_radius = math.ceil(math.sqrt(n)) n_sphere, n_0s = divmod(n, n_per_radius) diff --git a/tests/tools/gaussian_mixture/fit_gmm_pair_test.py b/tests/tools/gaussian_mixture/fit_gmm_pair_test.py index 8a2c7fdbf..5c776c9dd 100644 --- a/tests/tools/gaussian_mixture/fit_gmm_pair_test.py +++ b/tests/tools/gaussian_mixture/fit_gmm_pair_test.py @@ -11,7 +11,6 @@ # 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 typing import Tuple import pytest @@ -33,7 +32,7 @@ def samples( gmm_reference: gaussian_mixture.GaussianMixture, gmm_shifted: gaussian_mixture.GaussianMixture -) -> Tuple[jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array, jax.Array]: """Points drawn from the reference mixture and from its shifted variant.""" rng0, rng1 = jr.split(jr.key(0), 2) return ( @@ -48,7 +47,7 @@ class TestFitGmmPair: balanced=[False, True], weighted=[False, True], only_fast=0 ) def test_fit_gmm( - self, rng: jax.Array, samples: Tuple[jnp.ndarray, jnp.ndarray], + self, rng: jax.Array, samples: tuple[jax.Array, jax.Array], balanced: bool, weighted: bool ): # dumb integration test that makes sure nothing crashes diff --git a/tests/tools/gaussian_mixture/fit_gmm_test.py b/tests/tools/gaussian_mixture/fit_gmm_test.py index 2ded56af2..7a9bacda1 100644 --- a/tests/tools/gaussian_mixture/fit_gmm_test.py +++ b/tests/tools/gaussian_mixture/fit_gmm_test.py @@ -14,7 +14,6 @@ import pytest import jax -import jax.numpy as jnp import jax.random as jr import jax.test_util @@ -22,7 +21,7 @@ @pytest.fixture(scope="module") -def samples(gmm_reference: gaussian_mixture.GaussianMixture) -> jnp.ndarray: +def samples(gmm_reference: gaussian_mixture.GaussianMixture) -> jax.Array: """Points drawn from the reference mixture.""" return gmm_reference.sample(rng=jr.key(0), size=2000) @@ -30,7 +29,7 @@ def samples(gmm_reference: gaussian_mixture.GaussianMixture) -> jnp.ndarray: @pytest.mark.fast() class TestFitGmm: - def test_integration(self, rng: jax.Array, samples: jnp.ndarray): + def test_integration(self, rng: jax.Array, samples: jax.Array): # dumb integration test that makes sure nothing crashes # Fit a GMM to the samples diff --git a/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py b/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py index f451fd193..3b73b9f06 100644 --- a/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py +++ b/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py @@ -11,10 +11,10 @@ # 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 typing import Tuple import pytest +import chex import jax import jax.numpy as jnp import jax.random as jr @@ -23,7 +23,7 @@ from ott.tools.gaussian_mixture import gaussian_mixture, gaussian_mixture_pair -GMMPair_t = Tuple[gaussian_mixture.GaussianMixture, +GMMPair_t = tuple[gaussian_mixture.GaussianMixture, gaussian_mixture.GaussianMixture] N_COMPONENTS = 3 @@ -156,12 +156,12 @@ def test_flatten_unflatten( children, aux_data = jax.tree_util.tree_flatten(pair) pair_new = jax.tree_util.tree_unflatten(aux_data, children) - assert pair.gmm0 == pair_new.gmm0 - assert pair.gmm1 == pair_new.gmm1 + chex.assert_trees_all_equal(pair.gmm0, pair_new.gmm0) + chex.assert_trees_all_equal(pair.gmm1, pair_new.gmm1) assert pair.epsilon == pair_new.epsilon assert pair.tau == pair_new.tau assert pair.lock_gmm1 == pair_new.lock_gmm1 - assert pair == pair_new + chex.assert_trees_all_equal(pair, pair_new) @pytest.mark.fast.with_args( "epsilon,tau,lock_gmm1", diff --git a/tests/tools/gaussian_mixture/gaussian_mixture_test.py b/tests/tools/gaussian_mixture/gaussian_mixture_test.py index 71252a3f9..5334428c5 100644 --- a/tests/tools/gaussian_mixture/gaussian_mixture_test.py +++ b/tests/tools/gaussian_mixture/gaussian_mixture_test.py @@ -13,6 +13,7 @@ # limitations under the License. import pytest +import chex import jax import jax.numpy as jnp import jax.random as jr @@ -157,7 +158,7 @@ def test_flatten_unflatten(self, rng: jax.Array): children, aux_data = jax.tree_util.tree_flatten(gmm) gmm_new = jax.tree_util.tree_unflatten(aux_data, children) - assert gmm == gmm_new + chex.assert_trees_all_equal(gmm, gmm_new) def test_pytree_mapping(self, rng: jax.Array): gmm = gaussian_mixture.GaussianMixture.from_random( diff --git a/tests/tools/gaussian_mixture/gaussian_test.py b/tests/tools/gaussian_mixture/gaussian_test.py index 8f4baf21e..8bccb43bb 100644 --- a/tests/tools/gaussian_mixture/gaussian_test.py +++ b/tests/tools/gaussian_mixture/gaussian_test.py @@ -13,6 +13,7 @@ # limitations under the License. import pytest +import chex import jax import jax.numpy as jnp import jax.random as jr @@ -142,7 +143,7 @@ def test_flatten_unflatten(self, rng: jax.Array): children, aux_data = jtu.tree_flatten(g) g_new = jtu.tree_unflatten(aux_data, children) - assert g == g_new + chex.assert_trees_all_equal(g, g_new) def test_pytree_mapping(self, rng: jax.Array): g = gaussian.Gaussian.from_random(rng, n_dimensions=3) diff --git a/tests/tools/gaussian_mixture/probabilities_test.py b/tests/tools/gaussian_mixture/probabilities_test.py index 4cb0dc370..a0f794a3d 100644 --- a/tests/tools/gaussian_mixture/probabilities_test.py +++ b/tests/tools/gaussian_mixture/probabilities_test.py @@ -13,6 +13,7 @@ # limitations under the License. import pytest +import chex import jax import jax.numpy as jnp import jax.tree_util as jtu @@ -67,7 +68,7 @@ def test_flatten_unflatten(self): children, aux_data = jtu.tree_flatten(pp) pp_new = jtu.tree_unflatten(aux_data, children) np.testing.assert_array_equal(pp.params, pp_new.params) - assert pp == pp_new + chex.assert_trees_all_equal(pp, pp_new) def test_pytree_mapping(self): probs = jnp.array([0.1, 0.2, 0.3, 0.4]) diff --git a/tests/tools/gaussian_mixture/scale_tril_test.py b/tests/tools/gaussian_mixture/scale_tril_test.py index ff98bec90..dd6446b93 100644 --- a/tests/tools/gaussian_mixture/scale_tril_test.py +++ b/tests/tools/gaussian_mixture/scale_tril_test.py @@ -13,6 +13,7 @@ # limitations under the License. import pytest +import chex import jax import jax.numpy as jnp import jax.random as jr @@ -107,7 +108,7 @@ def test_flatten_unflatten(self, rng: jax.Array): children, aux_data = jtu.tree_flatten(scale) scale_new = jtu.tree_unflatten(aux_data, children) np.testing.assert_array_equal(scale.params, scale_new.params) - assert scale == scale_new + chex.assert_trees_all_equal(scale, scale_new) def test_pytree_mapping(self, rng: jax.Array): scale = scale_tril.ScaleTriL.from_random(rng=rng, n_dimensions=3) diff --git a/tests/tools/k_means_test.py b/tests/tools/k_means_test.py index b7841826f..477b4b1ad 100644 --- a/tests/tools/k_means_test.py +++ b/tests/tools/k_means_test.py @@ -11,7 +11,7 @@ # 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 typing import Any, Literal, Optional, Tuple, Union +from typing import Any, Literal import pytest @@ -29,9 +29,9 @@ def make_blobs( *args: Any, - cost_fn: Optional[Literal["sqeucl", "cosine"]] = None, + cost_fn: Literal["sqeucl", "cosine"] | None = None, **kwargs: Any -) -> Tuple[Union[jnp.ndarray, pointcloud.PointCloud], jnp.ndarray, jnp.ndarray]: +) -> tuple[jax.Array | pointcloud.PointCloud, jax.Array, jax.Array]: X, y, c = datasets.make_blobs(*args, return_centers=True, **kwargs) X, y, c = jnp.asarray(X), jnp.asarray(y), jnp.asarray(c) if cost_fn is None: @@ -47,10 +47,10 @@ def make_blobs( def compute_assignment( - x: jnp.ndarray, - centers: jnp.ndarray, - weights: Optional[jnp.ndarray] = None -) -> Tuple[jnp.ndarray, float]: + x: jax.Array, + centers: jax.Array, + weights: jax.Array | None = None +) -> tuple[jax.Array, float]: if weights is None: weights = jnp.ones(x.shape[0]) cost_matrix = pointcloud.PointCloud(x, centers).cost_matrix @@ -63,7 +63,7 @@ def compute_assignment( class TestKmeansPlusPlus: @pytest.mark.fast.with_args("n_local_trials", [None, 3], only_fast=-1) - def test_n_local_trials(self, rng: jax.Array, n_local_trials: Optional[int]): + def test_n_local_trials(self, rng: jax.Array, n_local_trials: int | None): n, k = 100, 4 rng1, rng2 = rng, jr.key(0) geom, _, c = make_blobs( @@ -104,7 +104,7 @@ def test_matches_sklearn(self, rng: jax.Array, k: int): def test_initialization_differentiable(self, rng: jax.Array): - def callback(x: jnp.ndarray) -> float: + def callback(x: jax.Array) -> float: geom = pointcloud.PointCloud(x) centers = k_means._k_means_plus_plus(geom, k=3, rng=jr.key(0)) _, inertia = compute_assignment(x, centers) @@ -335,7 +335,7 @@ def test_k_means_jitting( self, rng: jax.Array, init: Literal["k-means++", "random"] ): - def callback(x: jnp.ndarray) -> k_means.KMeansOutput: + def callback(x: jax.Array) -> k_means.KMeansOutput: return k_means.k_means( x, k=k, init=init, store_inner_errors=True, rng=jr.key(0) ) @@ -363,7 +363,7 @@ def test_k_means_differentiability( self, rng: jax.Array, jit: bool, force_scan: bool ): - def inertia(x: jnp.ndarray, w: jnp.ndarray) -> float: + def inertia(x: jax.Array, w: jax.Array) -> float: return k_means.k_means( x, k=k, diff --git a/tests/tools/sinkhorn_divergence_test.py b/tests/tools/sinkhorn_divergence_test.py index 54825f64f..36c742779 100644 --- a/tests/tools/sinkhorn_divergence_test.py +++ b/tests/tools/sinkhorn_divergence_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Dict, Optional, Tuple +from typing import Any import pytest @@ -399,7 +399,7 @@ def g(rng, n): # yapf: enable def test_euclidean_momentum_params( self, clouds: _utils.PointClouds, rng: jax.Array, - solve_kwargs: Dict[str, Any], epsilon: Optional[float] + solve_kwargs: dict[str, Any], epsilon: float | None ): # check if sinkhorn divergence solve_kwargs parameters used for # momentum/Anderson are properly overridden for the symmetric (x,x) and @@ -472,7 +472,7 @@ def test_gradient_generic_point_cloud_wrapper( ) @pytest.mark.parametrize("grid_size", [(5,), (2, 3), (3, 4, 5)]) - def test_grid_geometry(self, rng: jax.Array, grid_size: Tuple[int, ...]): + def test_grid_geometry(self, rng: jax.Array, grid_size: tuple[int, ...]): rng1, rng2 = jr.split(rng, 2) gs = (5,) diff --git a/tests/tools/sliced_test.py b/tests/tools/sliced_test.py index b4bfa615d..f898b24ee 100644 --- a/tests/tools/sliced_test.py +++ b/tests/tools/sliced_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Callable, Optional, Tuple +from collections.abc import Callable import pytest @@ -26,20 +26,17 @@ from ott.tools import sliced from tests import _utils -Projector = Callable[[jax.Array, jnp.ndarray], jnp.ndarray] +Projector = Callable[[jax.Array, jax.Array], jax.Array] -def custom_proj( - rng: jax.Array, x: jnp.ndarray, *, n_proj: int = 27 -) -> jnp.ndarray: +def custom_proj(rng: jax.Array, x: jax.Array, *, n_proj: int = 27) -> jax.Array: dim = x.shape[1] proj_m = jr.uniform(rng, (n_proj, dim)) return (x @ proj_m.T) ** 2 -def gen_data( - rng: jax.Array, n: int, m: int, dim: int -) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: +def gen_data(rng: jax.Array, n: int, m: int, + dim: int) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: c = _utils.random_clouds(rng, n=n, m=m, dim=dim, offset=0.0) return c.a, c.x, c.b, c.y @@ -49,8 +46,8 @@ class TestSliced: @pytest.mark.parametrize("proj_fn", [None, custom_proj]) @pytest.mark.parametrize("cost_fn", [costs.PNormP(1.3), None]) def test_random_projs( - self, rng: jax.Array, cost_fn: Optional[costs.CostFn], - proj_fn: Optional[Projector] + self, rng: jax.Array, cost_fn: costs.CostFn | None, + proj_fn: Projector | None ): n, m, dim, n_proj = 12, 17, 5, 13 rng_data, rng_w, rng_proj = jr.split(rng, 3) @@ -79,7 +76,7 @@ def test_random_projs( @pytest.mark.parametrize("cost_fn", [costs.SqPNorm(1.4), None]) def test_consistency_with_id( - self, rng: jax.Array, cost_fn: Optional[costs.CostFn] + self, rng: jax.Array, cost_fn: costs.CostFn | None ): n, m, dim = 11, 12, 4 _, x, _, y = gen_data(rng, n, m, dim) @@ -93,7 +90,7 @@ def test_consistency_with_id( np.testing.assert_allclose(out_lin, cost, rtol=1e-6, atol=1e-6) @pytest.mark.parametrize("proj_fn", [None, custom_proj]) - def test_diff(self, rng: jax.Array, proj_fn: Optional[Projector]): + def test_diff(self, rng: jax.Array, proj_fn: Projector | None): eps = 1e-4 n, m, dim = 13, 16, 7 rng_data, rng_dx = jr.split(rng, 2) diff --git a/tests/tools/soft_sort_test.py b/tests/tools/soft_sort_test.py index 31d8ffa2a..e2da063bd 100644 --- a/tests/tools/soft_sort_test.py +++ b/tests/tools/soft_sort_test.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Tuple import pytest @@ -29,7 +28,7 @@ class TestSoftSort: @pytest.mark.parametrize("shape", [(20,), (20, 1)]) - def test_sort_one_array(self, rng: jax.Array, shape: Tuple[int, ...]): + def test_sort_one_array(self, rng: jax.Array, shape: tuple[int, ...]): x = jr.uniform(rng, shape) xs = soft_sort.sort(x, axis=0) @@ -112,7 +111,7 @@ def test_multivariate_cdf_quantiles(self, rng: jax.Array): # Check passing custom sampler, must be still symmetric / centered on {.5}^d # Check passing custom epsilon also works. - def ball_sampler(k: jax.Array, s: Tuple[int, int]) -> jnp.ndarray: + def ball_sampler(k: jax.Array, s: tuple[int, int]) -> jax.Array: return 0.5 * (jr.ball(k, d=s[1], p=4, shape=(s[0],)) + 1.0) def mv_c_q(inputs, num_target_samples, rng, epsilon): @@ -288,7 +287,7 @@ def test_soft_sort_jacobian(self, rng: jax.Array, implicit: bool): z = jr.uniform(rngs[0], (b, n)) random_dir = jr.normal(rngs[1], (b,)) / b - def loss_fn(logits: jnp.ndarray) -> float: + def loss_fn(logits: jax.Array) -> float: im_d = None if implicit: # Ridge parameters are only used when using JAX's CG. diff --git a/tests/tools/unreg_test.py b/tests/tools/unreg_test.py index e5b24244e..d80c0d175 100644 --- a/tests/tools/unreg_test.py +++ b/tests/tools/unreg_test.py @@ -11,12 +11,10 @@ # 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 typing import Optional, Tuple import pytest import jax -import jax.numpy as jnp import numpy as np from ott.geometry import costs, pointcloud @@ -28,7 +26,7 @@ class TestHungarian: @pytest.mark.parametrize("cost_fn", [costs.PNormP(1.3), None]) - def test_matches_sink(self, rng: jax.Array, cost_fn: Optional[costs.CostFn]): + def test_matches_sink(self, rng: jax.Array, cost_fn: costs.CostFn | None): n, m, dim = 12, 12, 5 x, y = gen_data(rng, n, m, dim) geom = pointcloud.PointCloud(x, y, cost_fn=cost_fn, epsilon=.0005) @@ -52,6 +50,6 @@ def test_wass(self, rng: jax.Array, p: float): def gen_data(rng: jax.Array, n: int, m: int, - dim: int) -> Tuple[jnp.ndarray, jnp.ndarray]: + dim: int) -> tuple[jax.Array, jax.Array]: c = _utils.random_clouds(rng, n=n, m=m, dim=dim, offset=0.0) return c.x, c.y diff --git a/tests/utils_test.py b/tests/utils_test.py index 81bfae014..209e54843 100644 --- a/tests/utils_test.py +++ b/tests/utils_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import functools -from typing import Any, Optional +from typing import Any import pytest @@ -38,7 +38,7 @@ def test_batch_size(self, rng: jax.Array, batch_size: int): def test_pytree(self, rng: jax.Array): - def f(x: Any) -> jnp.ndarray: + def f(x: Any) -> jax.Array: return x["foo"]["bar"].std() + x["baz"].mean( ) + x["quux"][0] * x["quux"][1] @@ -61,7 +61,7 @@ def f(x: Any) -> jnp.ndarray: @pytest.mark.parametrize("in_axes", [0, 1, -1, -2, [0, None], (0, -2)]) def test_in_axes(self, rng: jax.Array, in_axes: Any, batch_size: int): - def f(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: + def f(x: jax.Array, y: jax.Array) -> jax.Array: x = jnp.atleast_2d(x) y = jnp.atleast_2d(y) return jnp.dot(x, y.T) @@ -90,7 +90,7 @@ def f(x: jnp.ndarray, y: jnp.ndarray) -> jnp.ndarray: ) def test_in_axes_pytree(self, rng: jax.Array, in_axes: Any): - def f(tree: Any) -> jnp.ndarray: + def f(tree: Any) -> jax.Array: x = tree[0]["foo"]["bar"] y = tree[0]["baz"] z, ((v,), w) = tree[1], tree[2] @@ -117,7 +117,7 @@ def f(tree: Any) -> jnp.ndarray: @pytest.mark.parametrize("out_axes", [0, 1, 2, -1, -2, -3]) def test_out_axes(self, rng: jax.Array, out_axes: int): - def f(x: jnp.ndarray, y: jnp.ndarray) -> Any: + def f(x: jax.Array, y: jax.Array) -> Any: return (x.sum() + y.sum()).reshape(1, 1) rng1, rng2 = jr.split(rng, 2) @@ -138,7 +138,7 @@ def f(x: jnp.ndarray, y: jnp.ndarray) -> Any: ) def test_out_axes_pytree(self, rng: jax.Array, out_axes: Any): - def f(x: jnp.ndarray) -> Any: + def f(x: jax.Array) -> Any: z = jnp.arange(9).reshape(3, 3) return x.mean(), {"x": {"y": jnp.ones(13)}}, (z,) @@ -157,7 +157,7 @@ def test_max_traces(self, rng: jax.Array, batch_size: int, n: int): @jax.jit @functools.partial(utils.batched_vmap, batch_size=batch_size) @chex.assert_max_traces(n=max_traces) - def fn(x: jnp.ndarray) -> jnp.ndarray: + def fn(x: jax.Array) -> jax.Array: return x.sum() chex.clear_trace_counter() @@ -206,7 +206,7 @@ def test_inconsistent_array_sizes(self, rng: jax.Array, batch_size: int): @pytest.mark.parametrize(("version", "msg"), [(None, "foo, bar, baz"), ("quux", None)]) -def test_deprecation_warning(version: Optional[str], msg: Optional[str]): +def test_deprecation_warning(version: str | None, msg: str | None): @utils.deprecate(version=version, alt=msg) def func() -> int: