Add backend specific TPU tests for DistributedEmbedding.

Under `keras_rs/src/layers/embedding`.

Author: hertschuh

.github/workflows/actions.yml

@@ -90,7 +90,7 @@ jobs:
       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
+      run: pytest keras_rs/src/layers/embedding/distributed_embedding_test.py keras_rs/src/layers/embedding/${{ matrix.backend }}
 
   check_format:
     name: Check the code format

keras_rs/src/layers/embedding/jax/distributed_embedding_test.py

@@ -338,7 +338,7 @@ def test_call(
 
         # Trigger layer.build(...) to initialize tables.
         sample_ids, sample_weights = keras_test_utils.create_random_samples(
-            feature_configs, ragged=ragged, seed=0
+            feature_configs, ragged=ragged, seed=0, max_ids_per_sample=10
         )
         inputs = layer.preprocess(sample_ids, sample_weights)
         _ = layer(inputs)