synopsis: fix some pyright warnings

synopsis.ipynb

@@ -44,7 +44,7 @@
    },
    "outputs": [],
    "source": [
-    "!pip install -q \"git+https://github.com/Astera-org/jax_loop_utils[tf-data]\" flax"
+    "%pip install -q \"git+https://github.com/Astera-org/jax_loop_utils[tf-data]\" flax"
    ]
   },
   {
@@ -72,7 +72,7 @@
     }
    ],
    "source": [
-    "!pip show jax-loop-utils"
+    "%pip show jax-loop-utils"
    ]
   },
   {
@@ -84,8 +84,7 @@
    "outputs": [],
    "source": [
     "import jax\n",
-    "import jax.numpy as jnp\n",
-    "import numpy as np"
+    "import jax.numpy as jnp"
    ]
   },
   {
@@ -208,15 +207,15 @@
    ],
    "source": [
     "import os\n",
+    "import pathlib\n",
     "from jax_loop_utils import metric_writers\n",
     "from jax_loop_utils.metric_writers.tf import SummaryWriter\n",
-    "import etils.epath as epath\n",
     "\n",
     "# Handy shortcut to create create async logging/tensorboard writer.\n",
     "\n",
     "\n",
     "def create_default_writer(\n",
-    "    logdir: epath.PathLike | None = None,\n",
+    "    logdir: os.PathLike | None = None,\n",
     "    *,\n",
     "    just_logging: bool = False,\n",
     "    asynchronous: bool = True,\n",
@@ -254,7 +253,7 @@
     "            )\n",
     "    writers = [metric_writers.LoggingWriter(collection=collection)]\n",
     "    if logdir is not None:\n",
-    "        logdir = epath.Path(logdir)\n",
+    "        logdir = pathlib.Path(logdir)\n",
     "        if collection is not None:\n",
     "            logdir /= collection\n",
     "        writers.append(SummaryWriter(os.fspath(logdir)))\n",
@@ -538,7 +537,7 @@
     "\n",
     "@flax.struct.dataclass  # <-- required for JAX transformations\n",
     "class MyMetrics(metrics.Collection):\n",
-    "    loss: metrics.Average.from_output(\"loss\")\n",
+    "    loss: metrics.Average.from_output(\"loss\")  # type: ignore[invalid-type-form]\n",
     "    accuracy: metrics.Accuracy\n",
     "\n",
     "\n",