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.py

@@ -1,5 +1,6 @@
 """JAX implementation of the TPU embedding layer."""
 
+import collections
 import math
 import typing
 from typing import Any, Mapping, Sequence, Union
@@ -445,6 +446,33 @@ def sparsecore_build(
         table_specs = embedding.get_table_specs(feature_specs)
         table_stacks = jte_table_stacking.get_table_stacks(table_specs)
 
+        # Compute max ids for stacked tables
+        stack_max_ids_per_partition: dict[str, int] = collections.defaultdict(
+            int
+        )
+        stack_max_unique_ids_per_partition: dict[str, int] = (
+            collections.defaultdict(int)
+        )
+
+        for stack_name, table_specs in table_stacks.items():
+            for table_spec in table_specs:
+                stack_max_ids_per_partition[stack_name] += (
+                    table_spec.max_ids_per_partition
+                )
+                stack_max_unique_ids_per_partition[stack_name] += (
+                    table_spec.max_unique_ids_per_partition
+                )
+
+        # stack name -> StackedTableSpec
+        stacked_table_specs = embedding.get_stacked_table_specs(feature_specs)
+        for stack_name, stacked_table_spec in stacked_table_specs.items():
+            stacked_table_spec.max_ids_per_partition = (
+                stack_max_ids_per_partition[stack_name]
+            )
+            stacked_table_spec.max_unique_ids_per_partition = (
+                stack_max_unique_ids_per_partition[stack_name]
+            )
+
         # Create variables for all stacked tables and slot variables.
         with sparsecore_distribution.scope():
             self._table_and_slot_variables = {

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

@@ -327,7 +327,10 @@ def test_call(
         jit: bool,
     ):
         table_configs = keras_test_utils.create_random_table_configs(
-            combiner=combiner, seed=10
+            combiner=combiner,
+            seed=10,
+            max_ids_per_partition=512,
+            max_unique_ids_per_partition=512,
         )
         feature_configs = keras_test_utils.create_random_feature_configs(
             table_configs=table_configs, seed=20