diff --git a/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb
index 5815822..0e6d3cf 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb
@@ -36,6 +36,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb b/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb
index 482043d..16edebf 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb
@@ -38,6 +38,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/evolving_in_small_worlds.ipynb b/self_organising_systems/biomakerca/examples/notebooks/evolving_in_small_worlds.ipynb
index 80cce25..c476780 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/evolving_in_small_worlds.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/evolving_in_small_worlds.ipynb
@@ -40,6 +40,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb
index f707b59..a828413 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb
@@ -37,6 +37,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb
index b0ae6cd..ee4c0a5 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb
@@ -36,6 +36,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb
index 3042e6e..ecfe0b9 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb
@@ -38,6 +38,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb
index fd93f06..e63bc0c 100644
--- a/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb
+++ b/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb
@@ -37,6 +37,29 @@
         "## Imports"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WpvHmJytfxpq"
+      },
+      "source": [
+        "We need to downgrade jax, unfortunately.\n",
+        "This is because we are using evojax which needs an old version of flax.\n",
+        "Old flax is not compatible with new jax ([see this discussion](https://github.com/google/flax/issues/3180))."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JgpxBfCFf7j2"
+      },
+      "outputs": [],
+      "source": [
+        "# downgrade jax with cuda.\n",
+        "!pip install \"jax[cuda12_pip]\"==0.4.23 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -74,7 +97,6 @@
         "from self_organising_systems.biomakerca.step_maker import step_env\n",
         "from self_organising_systems.biomakerca.display_utils import zoom, tile2d, add_text_to_img, imshow\n",
         "from self_organising_systems.biomakerca.custom_ipython_display import display\n",
-        "from self_organising_systems.biomakerca.env_logic import env_perform_multi_world_reproduce_update\n",
         "\n",
         "import cv2\n",
         "import numpy as np\n",