Remove workaround needed for evojax.

evojax now imports the latest flax, so we do not need to downgrade neither flax nor jax. PiperOrigin-RevId: 644286978
google-research · Jun 18, 2024 · 51de342 · 51de342
1 parent 8ff6c23
commit 51de342
Show file tree

Hide file tree

Showing 7 changed files with 5 additions and 187 deletions.
diff --git a/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/e2e_evolution.ipynb
@@ -36,29 +36,6 @@
         "## 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,
@@ -72,9 +49,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

diff --git a/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb b/self_organising_systems/biomakerca/examples/notebooks/eruption_example.ipynb
@@ -38,34 +38,11 @@
         "## 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,
       "metadata": {
-        "id": "287RU0sup5J2"
+        "id": "CnTgss2gWcHY"
       },
       "outputs": [],
       "source": [
@@ -74,9 +51,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

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
@@ -40,34 +40,11 @@
         "## 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,
       "metadata": {
-        "id": "287RU0sup5J2"
+        "id": "bVM4KsV0Wise"
       },
       "outputs": [],
       "source": [
@@ -76,9 +53,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

diff --git a/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/interactive_evolution.ipynb
@@ -37,34 +37,11 @@
         "## 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,
       "metadata": {
-        "id": "287RU0sup5J2"
+        "id": "_FW4_grEWn4f"
       },
       "outputs": [],
       "source": [
@@ -73,9 +50,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

diff --git a/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb b/self_organising_systems/biomakerca/examples/notebooks/petri_evolution.ipynb
@@ -36,34 +36,11 @@
         "## 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,
       "metadata": {
-        "id": "287RU0sup5J2"
+        "id": "X_3-KqY3WtXp"
       },
       "outputs": [],
       "source": [
@@ -72,9 +49,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_configuration.ipynb
@@ -38,34 +38,11 @@
         "## 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,
       "metadata": {
-        "id": "287RU0sup5J2"
+        "id": "u0WhscjnW0FO"
       },
       "outputs": [],
       "source": [
@@ -74,9 +51,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/[email protected]#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },

diff --git a/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb b/self_organising_systems/biomakerca/examples/notebooks/run_sexual_configuration.ipynb
@@ -37,29 +37,6 @@
         "## 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,
@@ -73,9 +50,6 @@
         "!pip install --upgrade -e git+https://github.com/google-research/self-organising-systems.git#egg=self_organising_systems\u0026subdirectory=biomakerca\n",
         "# activate the locally installed package (otherwise a runtime restart is required)\n",
         "import pkg_resources\n",
-        "import importlib\n",
-        "# Reload the resources because we uninstalled and reinstalled some packages.\n",
-        "importlib.reload(pkg_resources)\n",
         "pkg_resources.get_distribution(\"self_organising_systems\").activate()"
       ]
     },