DistributedEmbedding for JAX example needs more recent JAX version. (#2179)

The preinstalled version of JAX on colab no longer works with the latest version of the `jax-tpu-embedding`.

Author: hertschuh

examples/keras_rs/distributed_embedding_jax.py

@@ -24,6 +24,7 @@
 """
 
 """shell
+pip install -q -U jax[tpu]>=0.7.0
 pip install -q jax-tpu-embedding
 pip install -q tensorflow-cpu
 pip install -q keras-rs

examples/keras_rs/ipynb/distributed_embedding_jax.ipynb

@@ -43,6 +43,7 @@
    },
    "outputs": [],
    "source": [
+    "!pip install -q -U jax[tpu]>=0.7.0\n",
     "!pip install -q jax-tpu-embedding\n",
     "!pip install -q tensorflow-cpu\n",
     "!pip install -q keras-rs"

examples/keras_rs/md/distributed_embedding_jax.md

@@ -27,6 +27,7 @@ libraries.
 
 
 ```python
+!pip install -q -U jax[tpu]>=0.7.0
 !pip install -q jax-tpu-embedding
 !pip install -q tensorflow-cpu
 !pip install -q keras-rs