keras-rs-nightly 0.2.2.dev202508050345__tar.gz → 0.2.2.dev202508070344__tar.gz

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202508050345
+Version: 0.2.2.dev202508070344
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -146,14 +146,14 @@ class DistributedEmbedding(keras.layers.Layer):
     feature1 = keras_rs.layers.FeatureConfig(
         name="feature1",
         table=table1,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
     )
     feature2 = keras_rs.layers.FeatureConfig(
         name="feature2",
         table=table2,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
     )
 
     feature_configs = {

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs/src/layers/embedding/distributed_embedding_config.py

@@ -102,7 +102,10 @@ class FeatureConfig:
         input_shape: The input shape of the feature. The feature fed into the
             layer has to match the shape. Note that for ragged dimensions in the
             input, the dimension provided here presents the maximum value;
-            anything larger will be truncated.
+            anything larger will be truncated. Also note that the first
+            dimension represents the global batch size. For example, on TPU,
+            this represents the total number of samples that are dispatched to
+            all the TPUs connected to the current host.
         output_shape: The output shape of the feature activation. What is
             returned by the embedding layer has to match this shape.
     """

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs/src/layers/embedding/tensorflow/config_conversion.py

@@ -56,6 +56,7 @@ OPTIMIZER_MAPPINGS = {
 def translate_keras_rs_configuration(
     feature_configs: types.Nested[FeatureConfig],
     table_stacking: str | Sequence[str] | Sequence[Sequence[str]],
+    num_replicas_in_sync: int,
 ) -> tuple[
     types.Nested[tf.tpu.experimental.embedding.FeatureConfig],
     tf.tpu.experimental.embedding.SparseCoreEmbeddingConfig,
@@ -72,7 +73,10 @@ def translate_keras_rs_configuration(
     """
     tables: dict[TableConfig, tf.tpu.experimental.embedding.TableConfig] = {}
     feature_configs = keras.tree.map_structure(
-        lambda f: translate_keras_rs_feature_config(f, tables), feature_configs
+        lambda f: translate_keras_rs_feature_config(
+            f, tables, num_replicas_in_sync
+        ),
+        feature_configs,
     )
 
     # max_ids_per_chip_per_sample
@@ -107,6 +111,7 @@ def translate_keras_rs_configuration(
 def translate_keras_rs_feature_config(
     feature_config: FeatureConfig,
     tables: dict[TableConfig, tf.tpu.experimental.embedding.TableConfig],
+    num_replicas_in_sync: int,
 ) -> tf.tpu.experimental.embedding.FeatureConfig:
     """Translates a Keras RS feature config to a TensorFlow TPU feature config.
 
@@ -120,18 +125,46 @@ def translate_keras_rs_feature_config(
     Returns:
       The TensorFlow TPU feature config.
     """
+    if num_replicas_in_sync <= 0:
+        raise ValueError(
+            "`num_replicas_in_sync` must be positive, "
+            f"but got {num_replicas_in_sync}."
+        )
+
     table = tables.get(feature_config.table, None)
     if table is None:
         table = translate_keras_rs_table_config(feature_config.table)
         tables[feature_config.table] = table
 
+    if len(feature_config.output_shape) < 2:
+        raise ValueError(
+            f"Invalid `output_shape` {feature_config.output_shape} in "
+            f"`FeatureConfig` {feature_config}. It must have at least 2 "
+            "dimensions: a batch dimension and an embedding dimension."
+        )
+
+    # Exclude last dimension, TensorFlow's TPUEmbedding doesn't want it.
+    output_shape = list(feature_config.output_shape[0:-1])
+
+    batch_size = output_shape[0]
+    per_replica_batch_size: int | None = None
+    if batch_size is not None:
+        if batch_size % num_replicas_in_sync != 0:
+            raise ValueError(
+                f"Invalid `output_shape` {feature_config.output_shape} in "
+                f"`FeatureConfig` {feature_config}. Batch size {batch_size} is "
+                f"not a multiple of the number of TPUs {num_replicas_in_sync}."
+            )
+        per_replica_batch_size = batch_size // num_replicas_in_sync
+
+    # TensorFlow's TPUEmbedding wants the per replica batch size.
+    output_shape = [per_replica_batch_size] + output_shape[1:]
+
     # max_sequence_length
     return tf.tpu.experimental.embedding.FeatureConfig(
         name=feature_config.name,
         table=table,
-        output_shape=feature_config.output_shape[
-            0:-1
-        ],  # exclude last dimension
+        output_shape=output_shape,
     )
 
 

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py

@@ -107,7 +107,9 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
                 )
             self._tpu_feature_configs, self._sparse_core_embedding_config = (
                 config_conversion.translate_keras_rs_configuration(
-                    feature_configs, table_stacking
+                    feature_configs,
+                    table_stacking,
+                    strategy.num_replicas_in_sync,
                 )
             )
             if tpu_embedding_feature == EMBEDDING_FEATURE_V1:

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs/src/version.py

@@ -1,7 +1,7 @@
 from keras_rs.src.api_export import keras_rs_export
 
 # Unique source of truth for the version number.
-__version__ = "0.2.2.dev202508050345"
+__version__ = "0.2.2.dev202508070344"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.2.2.dev202508050345 → keras_rs_nightly-0.2.2.dev202508070344}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202508050345
+Version: 0.2.2.dev202508070344
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0