Fix missing jax import in TessellateIPU notebook.

graphcore-research · Sep 27, 2023 · 5776225 · 5776225
1 parent 1ca7c92
commit 5776225
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/notebooks/01-tessellate-ipu-tile-api-basics.ipynb b/notebooks/01-tessellate-ipu-tile-api-basics.ipynb
@@ -200,6 +200,7 @@
     }
    ],
    "source": [
+    "import jax\n",
     "import numpy as np\n",
     "\n",
     "from tessellate_ipu import tile_put_replicated, tile_put_sharded\n",