diff --git a/notebooks/01-tessellate-ipu-tile-api-basics.ipynb b/notebooks/01-tessellate-ipu-tile-api-basics.ipynb
index 0870c1f..df4c716 100644
--- 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",