Misc/update typing

.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

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"

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

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

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