keras-rs-nightly 0.2.2.dev202508060343__py3-none-any.whl → 0.2.2.dev202508070344__py3-none-any.whl

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/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/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.dev202508060343"
+__version__ = "0.2.2.dev202508070344"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.2.2.dev202508060343.dist-info → keras_rs_nightly-0.2.2.dev202508070344.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202508060343
+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.dev202508060343.dist-info → keras_rs_nightly-0.2.2.dev202508070344.dist-info}/RECORD

@@ -5,7 +5,7 @@ keras_rs/metrics/__init__.py,sha256=Qxpf6OFooIL9TIn2l3WgOea3HFRG0hq02glPAxtMZ9c,
 keras_rs/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/api_export.py,sha256=RsmG-DvO-cdFeAF9W6LRzms0kvtm-Yp9BAA_d-952zI,510
 keras_rs/src/types.py,sha256=1A-oLRdX1-f2DsVZBcNl8qNsaH8pM-gnleLT9FWZWBw,1189
-keras_rs/src/version.py,sha256=THPSXnX6_BLgHCwK3ZXEGVqhwtxAiOP-YoUSJnuAwEE,224
+keras_rs/src/version.py,sha256=u9dAyD9wPmdZRGUcU0EBtAFVCmSHpqmIJPDURbAXDoo,224
 keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/embedding/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=J1mW7sgOgfkR8s-n3uyuqMb_iwPgrD6j0J5OBSfXQt8,45006
@@ -19,8 +19,8 @@ keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=jLqEuh_7hHM4ba
 keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=8LigXjPr7uQaUOdZM6yoLGoPYdRcbkXkFeL_sJoQ6uQ,8223
 keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=EHrQjPLl94STLWf9g8Ew8nuwupXRq-a_QmvFlXV6G6A,20331
 keras_rs/src/layers/embedding/tensorflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/embedding/tensorflow/config_conversion.py,sha256=i0iNXHi9E2UNh32Q5_kQtB9W5Cbcbm-tbWwRS1FHrfE,11589
-keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=WPkRLJ7Tz272A0mMb5LCHlnNMYl_mTF2VxKpObcfQic,17206
+keras_rs/src/layers/embedding/tensorflow/config_conversion.py,sha256=6aU7B-SX-WUqH8UC1amXg6BpFdalFmbr7rYQiiH_k4A,12883
+keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=-K6l3PpIqOEVCzcVZN_cJMx9ddJnRKBTj486CP3wkps,17278
 keras_rs/src/layers/feature_interaction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/feature_interaction/dot_interaction.py,sha256=Rs8xIHXNWQNiwjp_xzvQRmTSV1AyhJjDgVc3K5pTmrQ,8530
 keras_rs/src/layers/feature_interaction/feature_cross.py,sha256=Wq_eQvO0WTRlep69QbKi8TwY8bnFoF9vreP_j6ZHNFE,8666
@@ -50,7 +50,7 @@ keras_rs/src/metrics/utils.py,sha256=fGTo8j0ykVE5Y3yQCS2orSFcHY20Uxt0NazyPsybUsw
 keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/utils/doc_string_utils.py,sha256=CmqomepmaYcvpACpXEXkrJb8DMnvIgmYK-lJ53lYarY,1675
 keras_rs/src/utils/keras_utils.py,sha256=dc-NFzs3a-qmRw0vBDiMslPLfrm9yymGduLWesXPhuY,2123
-keras_rs_nightly-0.2.2.dev202508060343.dist-info/METADATA,sha256=kpIe9X61t56x-jdjn7jOwgazH7OMHKSBt5MwKrMml9s,5273
-keras_rs_nightly-0.2.2.dev202508060343.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.2.2.dev202508060343.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.2.2.dev202508060343.dist-info/RECORD,,
+keras_rs_nightly-0.2.2.dev202508070344.dist-info/METADATA,sha256=rBjnDQocZX2HP0E7Rj72zR7icvhp_ziuzH46k1ESB20,5273
+keras_rs_nightly-0.2.2.dev202508070344.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.2.2.dev202508070344.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.2.2.dev202508070344.dist-info/RECORD,,