DistributedEmbedding: do not hardcode the number of SparseCores with TF. (keras-team#147)

The number of SparseCore chips per TPU is now retrieved from the strategy. This makes the `distributed_embedding_tests.py` pass on V6e.

Author: hertschuh

keras_rs/src/layers/embedding/distributed_embedding_test.py

@@ -611,7 +611,9 @@ def test_correctness(
 
         self.assertEqual(res.shape, (self.batch_size, EMBEDDING_OUTPUT_DIM))
 
-        tables = layer.get_embedding_tables()
+        with self._strategy.scope():
+            tables = layer.get_embedding_tables()
+
         emb = tables["table"]
 
         if input_type == "dense":

keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py

@@ -281,8 +281,18 @@ def _sparsecore_call(
     def _sparsecore_get_embedding_tables(self) -> dict[str, types.Tensor]:
         tables: dict[str, types.Tensor] = {}
         strategy = tf.distribute.get_strategy()
-        # 4 is the number of sparsecores per chip
-        num_shards = strategy.num_replicas_in_sync * 4
+        if not self._is_tpu_strategy(strategy):
+            raise RuntimeError(
+                "`DistributedEmbedding.get_embedding_tables` needs to be "
+                "called under the TPUStrategy that DistributedEmbedding was "
+                f"created with, but is being called under strategy {strategy}. "
+                "Please use `with strategy.scope()` when calling "
+                "`get_embedding_tables`."
+            )
+
+        tpu_hardware = strategy.extended.tpu_hardware_feature
+        num_sc_per_device = tpu_hardware.num_embedding_devices_per_chip
+        num_shards = strategy.num_replicas_in_sync * num_sc_per_device
 
         def populate_table(
             feature_config: tf.tpu.experimental.embedding.FeatureConfig,