Tweaks to DistributedEmbedding for JAX example. (#2171)

- Add `pip install`s
- Cast counts to int to not use `np.array`s in configuration

Author: hertschuh

examples/keras_rs/distributed_embedding_jax.py

@@ -23,6 +23,12 @@
 libraries.
 """
 
+"""shell
+pip install -q jax-tpu-embedding
+pip install -q tensorflow-cpu
+pip install -q keras-rs
+"""
+
 import os
 
 os.environ["KERAS_BACKEND"] = "jax"
@@ -64,7 +70,7 @@
 index in the user embedding table.
 """
 
-users_count = (
+users_count = int(
     ratings.map(lambda x: tf.strings.to_number(x["user_id"], out_type=tf.int32))
     .reduce(tf.constant(0, tf.int32), tf.maximum)
     .numpy()
@@ -75,7 +81,7 @@
 as an index in the movie embedding table.
 """
 
-movies_count = movies.cardinality().numpy()
+movies_count = int(movies.cardinality().numpy())
 
 """
 The inputs to the model are the user IDs and movie IDs and the labels are the

examples/keras_rs/ipynb/distributed_embedding_jax.ipynb

@@ -35,6 +35,19 @@
     "libraries."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install -q jax-tpu-embedding\n",
+    "!pip install -q tensorflow-cpu\n",
+    "!pip install -q keras-rs"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 0,
@@ -126,7 +139,7 @@
    },
    "outputs": [],
    "source": [
-    "users_count = (\n",
+    "users_count = int(\n",
     "    ratings.map(lambda x: tf.strings.to_number(x[\"user_id\"], out_type=tf.int32))\n",
     "    .reduce(tf.constant(0, tf.int32), tf.maximum)\n",
     "    .numpy()\n",
@@ -151,7 +164,7 @@
    },
    "outputs": [],
    "source": [
-    "movies_count = movies.cardinality().numpy()"
+    "movies_count = int(movies.cardinality().numpy())"
    ]
   },
   {

examples/keras_rs/md/distributed_embedding_jax.md

@@ -26,6 +26,12 @@ Let's begin by choosing JAX as the backend and importing all the necessary
 libraries.
 
 
+```python
+!pip install -q jax-tpu-embedding
+!pip install -q tensorflow-cpu
+!pip install -q keras-rs
+```
+
 ```python
 import os
 
@@ -73,7 +79,7 @@ index in the user embedding table.
 
 
 ```python
-users_count = (
+users_count = int(
     ratings.map(lambda x: tf.strings.to_number(x["user_id"], out_type=tf.int32))
     .reduce(tf.constant(0, tf.int32), tf.maximum)
     .numpy()
@@ -85,7 +91,7 @@ as an index in the movie embedding table.
 
 
 ```python
-movies_count = movies.cardinality().numpy()
+movies_count = int(movies.cardinality().numpy())
 ```
 
 The inputs to the model are the user IDs and movie IDs and the labels are the