From 577622524b47d9832a4b683f223623bbc6ef6c7a Mon Sep 17 00:00:00 2001
From: Paul Balanca <paulb@graphcore.ai>
Date: Wed, 27 Sep 2023 15:30:26 +0100
Subject: [PATCH] Fix missing jax import in TessellateIPU notebook.

---
 notebooks/01-tessellate-ipu-tile-api-basics.ipynb | 1 +
 1 file changed, 1 insertion(+)

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