Add TPU tests for JAX and Tensorflow. (#160)

This is using the self-hosted TPU runners.

Only the multi-backend `keras_rs/src/layers/embedding/distributed_embedding_test.py` test is run for now. Many other tests have failures. They will be addressed in subsequent PRs.

- Modified `distributed_embedding_test.py` to replace the `TPU` flag with a `TPU_NAME` environment variable as plumbing flag through `pytest` is needlessly complicated.
- Modified `distributed_embedding_test.py` to not require JAX to be installed when running against the TensorFlow backend.

Author: hertschuh

.github/workflows/actions.yml

@@ -45,6 +45,53 @@ jobs:
     - name: Test with pytest
       run: |
         pytest keras_rs/
+
+  run_tests_in_container:
+    name: Test the code on TPU
+    runs-on: linux-x86-ct6e-44-1tpu
+
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [tensorflow, jax]
+
+    container:
+      image: python:3.11-slim
+      options: --privileged --network host
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Install Dependencies
+      run: |
+        pip install --no-cache-dir -U pip && \
+        pip install --no-cache-dir -r requirements-${{ matrix.backend }}-tpu.txt
+    
+    - name: Set Keras Backend
+      run: |
+        echo "KERAS_BACKEND=${{ matrix.backend }}" >> $GITHUB_ENV
+        echo "TPU_NAME=local" >> $GITHUB_ENV
+
+    - name: Set TF Specific Environment Variables
+      if: ${{ matrix.backend == 'tensorflow'}}
+      run: |
+        echo "PJRT_DEVICE=TPU" >> $GITHUB_ENV
+        echo "NEXT_PLUGGABLE_DEVICE_USE_C_API=true" >> $GITHUB_ENV
+        echo "TF_XLA_FLAGS=--tf_mlir_enable_mlir_bridge=true" >> $GITHUB_ENV
+        pip show libtpu | grep "^Location: " | sed "s/^Location: \(.*\)$/TF_PLUGGABLE_DEVICE_LIBRARY_PATH=\1\/libtpu\/libtpu.so/1" >> $GITHUB_ENV
+
+    - name: Verify TF Installation
+      if: ${{ matrix.backend == 'tensorflow'}}
+      run: python3 -c "import tensorflow as tf; print('Tensorflow devices:', tf.config.list_logical_devices())"
+
+    - name: Verify JAX Installation
+      if: ${{ matrix.backend == 'jax'}}
+      run: python3 -c "import jax; print('JAX devices:', jax.devices())"
+
+    - name: Test with pytest
+      run: pytest keras_rs/src/layers/embedding/distributed_embedding_test.py
+
   check_format:
     name: Check the code format
     runs-on: ubuntu-latest

keras_rs/src/layers/embedding/distributed_embedding_test.py

@@ -4,13 +4,9 @@
 import os
 import tempfile
 
-import jax
-import jax.experimental.sparse as jax_sparse
-import jax.numpy as jnp
 import keras
 import numpy as np
 import tensorflow as tf
-from absl import flags
 from absl.testing import absltest
 from absl.testing import parameterized
 
@@ -19,8 +15,12 @@
 from keras_rs.src.layers.embedding import distributed_embedding
 from keras_rs.src.layers.embedding import distributed_embedding_config as config
 
-FLAGS = flags.FLAGS
-_TPU = flags.DEFINE_string("tpu", None, "The TPU to use for TPUStrategy.")
+try:
+    import jax
+    import jax.experimental.sparse as jax_sparse
+except ImportError:
+    jax = None
+    jax_sparse = None
 
 
 FEATURE1_EMBEDDING_OUTPUT_DIM = 7
@@ -50,29 +50,32 @@ def experimental_distribute_dataset(self, dataset, options=None):
 class JaxDummyStrategy(DummyStrategy):
     @property
     def num_replicas_in_sync(self):
-        return len(jax.devices("tpu"))
+        return jax.device_count("tpu")
+
+
+def ragged_bool_true(self):
+    return True
 
 
 class DistributedEmbeddingTest(testing.TestCase, parameterized.TestCase):
     def setUp(self):
         super().setUp()
-        try:
-            self.on_tpu = _TPU.value is not None
-        except flags.UnparsedFlagAccessError:
-            self.on_tpu = False
-
+        self.on_tpu = "TPU_NAME" in os.environ
         self.placement = "sparsecore" if self.on_tpu else "default_device"
 
         if keras.backend.backend() == "tensorflow":
             tf.debugging.disable_traceback_filtering()
 
         if keras.backend.backend() == "tensorflow" and self.on_tpu:
+            # Workaround for a bug preventing weights from being ragged tensors.
+            # The fix in TensorFlow was added after 2.19.1:
+            # https://github.com/tensorflow/tensorflow/commit/185f2f58bafc6410125080264d5d7730e1fa1eb2
+            tf.RaggedTensor.__bool__ = ragged_bool_true
+
             # FLAGS.xla_sparse_core_max_ids_per_partition_per_sample = 16
             # FLAGS.xla_sparse_core_max_unique_ids_per_partition_per_sample = 16
 
-            resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
-                tpu=_TPU.value
-            )
+            resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
             tf.config.experimental_connect_to_cluster(resolver)
 
             topology = tf.tpu.experimental.initialize_tpu_system(resolver)
@@ -187,7 +190,10 @@ def create_tensor(feature_config, op):
                 sequence_length = feature_config.input_shape[-1]
                 indices = [[i, i % sequence_length] for i in range(batch_size)]
                 return jax_sparse.BCOO(
-                    (jnp.asarray(op((batch_size,))), jnp.asarray(indices)),
+                    (
+                        jax.numpy.asarray(op((batch_size,))),
+                        jax.numpy.asarray(indices),
+                    ),
                     shape=(batch_size, sequence_length),
                     unique_indices=True,
                 )
@@ -534,16 +540,18 @@ def test_correctness(
             elif keras.backend.backend() == "jax":
                 inputs = jax_sparse.BCOO(
                     (
-                        jnp.asarray([1, 2, 3, 4, 5] * num_repeats),
-                        jnp.asarray(indices),
+                        jax.numpy.asarray([1, 2, 3, 4, 5] * num_repeats),
+                        jax.numpy.asarray(indices),
                     ),
                     shape=(self.batch_size, 4),
                     unique_indices=True,
                 )
                 weights = jax_sparse.BCOO(
                     (
-                        jnp.asarray([1.0, 1.0, 2.0, 3.0, 4.0] * num_repeats),
-                        jnp.asarray(indices),
+                        jax.numpy.asarray(
+                            [1.0, 1.0, 2.0, 3.0, 4.0] * num_repeats
+                        ),
+                        jax.numpy.asarray(indices),
                     ),
                     shape=(self.batch_size, 4),
                     unique_indices=True,

requirements-jax-tpu.txt

@@ -0,0 +1,10 @@
+# Tensorflow cpu-only version.
+tensorflow-cpu>=2.20.0
+
+# Jax with TPU support.
+jax[tpu]
+
+# Support for TPU embeddings.
+jax-tpu-embedding
+
+-r requirements-common.txt

requirements-tensorflow-tpu.txt

@@ -0,0 +1,4 @@
+# Tensorflow with TPU support.
+tensorflow-tpu==2.19.1
+
+-r requirements-common.txt