Skip to content

Commit

Permalink
Downgrade jax to 0.4.23 in biomakerca notebooks.
Browse files Browse the repository at this point in the history
This is because we are using evojax which needs an old version of flax. Old flax is not compatible with new jax ([see this discussion](google/flax#3180)).

This is a temporary patch. Hopefully. We will figure out how to either fix evojax or stop using it for future experiments.

Also, fix a bug on the sexual reproduction colab where I was importing a function that I deleted, causing the colab to fail.

PiperOrigin-RevId: 638631920
  • Loading branch information
oteret authored and Selforg Gardener committed May 30, 2024
1 parent f188d30 commit 8ff6c23
Show file tree
Hide file tree
Showing 7 changed files with 161 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down

0 comments on commit 8ff6c23

Please sign in to comment.