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

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/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/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.dev202508060343"
 
 
 @keras_rs_export("keras_rs.version")

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

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

@@ -5,12 +5,12 @@ 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=M8Lcch9FbrXIn-Syj211qspj2dEDiyksD18uOmw8-AM,224
+keras_rs/src/version.py,sha256=THPSXnX6_BLgHCwK3ZXEGVqhwtxAiOP-YoUSJnuAwEE,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=Hr54K8Y3OpUQc7ExtftLONoY7w4x4R08BkkMNtrpbZw,45026
+keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=J1mW7sgOgfkR8s-n3uyuqMb_iwPgrD6j0J5OBSfXQt8,45006
 keras_rs/src/layers/embedding/distributed_embedding.py,sha256=94jxUHoGK3Gs9yfV0KxFTuqPo7XFnhgCNlO2FEeiSgM,1072
-keras_rs/src/layers/embedding/distributed_embedding_config.py,sha256=AWPmZBir1shhqNP6U_jiQ9lsBhMXVikW4B5VnzLsvPg,5579
+keras_rs/src/layers/embedding/distributed_embedding_config.py,sha256=nmDqZlXiZGxc2UmEaUULaOvRjwq427ReSltDGEBrv5I,5816
 keras_rs/src/layers/embedding/embed_reduce.py,sha256=c-MnEw1-KWs0jTf0JJ_ZBOY-9hRkiFyu989Dof3DnS8,12343
 keras_rs/src/layers/embedding/jax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/embedding/jax/checkpoint_utils.py,sha256=wZ4I5WZVNg5WnrD2j7nhAXgLzDc7xMrUEkSAOx5Sz5c,3495
@@ -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.dev202508050345.dist-info/METADATA,sha256=ziJ6IdTpbj--g7l4hlNVsozpxk9I-Pr7iTs_APV3vxw,5273
-keras_rs_nightly-0.2.2.dev202508050345.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.2.2.dev202508050345.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.2.2.dev202508050345.dist-info/RECORD,,
+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,,