keras-rs-nightly 0.3.1.dev202512130338__py3-none-any.whl → 0.4.1.dev202601250348__py3-none-any.whl

keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -457,6 +457,10 @@ class DistributedEmbedding(keras.layers.Layer):
             tables in the inner lists together. Note that table stacking is not
             supported on older TPUs, in which case the default value of `"auto"`
             will be interpreted as no table stacking.
+        update_stats: If True, `'max_ids_per_partition'`,
+            `'max_unique_ids_per_partition'` and
+            `'suggested_coo_buffer_size_per_device'` are updated during
+            training. This argument can be set to True only for the JAX backend.
         **kwargs: Additional arguments to pass to the layer base class.
     """
 
@@ -467,6 +471,7 @@ class DistributedEmbedding(keras.layers.Layer):
         table_stacking: (
             str | Sequence[str] | Sequence[Sequence[str]]
         ) = "auto",
+        update_stats: bool = False,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -486,6 +491,8 @@ class DistributedEmbedding(keras.layers.Layer):
                 table_stacking,
             )
 
+        self.update_stats = update_stats
+
     @keras_utils.no_automatic_dependency_tracking
     def _init_feature_configs_structures(
         self,

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

@@ -407,7 +407,9 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             if isinstance(table_stacking, str):
                 if table_stacking == "auto":
                     jte_table_stacking.auto_stack_tables(
-                        feature_specs, global_device_count, num_sc_per_device
+                        feature_specs,
+                        global_device_count,
+                        num_sc_per_device,
                     )
                 else:
                     raise ValueError(
@@ -644,7 +646,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             num_sc_per_device,
         )
 
-        if training:
+        if training and self.update_stats:
             # Synchronize input statistics across all devices and update the
             # underlying stacked tables specs in the feature specs.
 

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

@@ -35,8 +35,15 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
         table_stacking: (
             str | Sequence[str] | Sequence[Sequence[str]]
         ) = "auto",
+        update_stats: bool = False,
         **kwargs: Any,
     ) -> None:
+        # `update_stats` is supported only on JAX.
+        if update_stats:
+            raise ValueError(
+                "`update_stats` cannot be True for the TensorFlow backend."
+            )
+
         # Intercept arguments that are supported only on TensorFlow.
         self._optimizer = kwargs.pop("optimizer", None)
         self._pipeline_execution_with_tensor_core = kwargs.pop(

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.3.1.dev202512130338"
+__version__ = "0.4.1.dev202601250348"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.3.1.dev202512130338.dist-info → keras_rs_nightly-0.4.1.dev202601250348.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202512130338
+Version: 0.4.1.dev202601250348
 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.3.1.dev202512130338.dist-info → keras_rs_nightly-0.4.1.dev202601250348.dist-info}/RECORD

@@ -5,22 +5,22 @@ 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=E7EoQd2LJhzxrYtMh6GLQgtbm2hzCwUNr9AcZw0MOgc,224
+keras_rs/src/version.py,sha256=7F19b6JBtXTYKOxUk6K2Y_nS6JGidlVq305CdH3935o,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=RkXZ6notj3Cq6ryR9w30Wb8UlaWjLcUK2Os9ZUQvuhY,45568
+keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=REpNKTvkS4eoAi9t1DohnDMfWwMgfxDN4ByWji3aALM,45906
 keras_rs/src/layers/embedding/distributed_embedding.py,sha256=94jxUHoGK3Gs9yfV0KxFTuqPo7XFnhgCNlO2FEeiSgM,1072
 keras_rs/src/layers/embedding/distributed_embedding_config.py,sha256=L41x6W1xcXI-3m94nOh_OsHn6OYpoynakKJvNboJnvE,5762
 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
 keras_rs/src/layers/embedding/jax/config_conversion.py,sha256=Di1UzRwLgGHd7RuWYJMj2mCOr1u9MseFEWaYKnwD9Bs,16742
-keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=lY7cwJgVUDpYUH-n7AcA1qlnSzhtckNtD3UqtpAYXz8,32267
+keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=D2tJehcy10w30KzYKnnV30WwCUvgLUm2YQM3Twwge9M,32338
 keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=a90tWTbU9tkFdESG3xir9PTtcvb1cmYR8vl5dDK9PSY,8703
 keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=5rQGli4Qflg0BU2-j_-4xbBxSqopqbtjkY2KKYWq64Y,7329
 keras_rs/src/layers/embedding/tensorflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/embedding/tensorflow/config_conversion.py,sha256=HpuDthRQQ3X0EO8dW6OAdcgTODkujZlx_swgreVwXyk,13220
-keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=TBPYV8gP3ZnAFEwtxmWr_Rp3s-Cj0RrKzF6UOALJ4B0,17817
+keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=rkxnPzMHmq82FEzrLrO13NhDHPiX-3PxRM3AUE6Rv10,18050
 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
@@ -52,7 +52,7 @@ keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuF
 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/src/utils/tpu_test_utils.py,sha256=mQVBrI-CCBbXwQxBq1yDKGMwYhm4g4k3_AaSy_sCs0U,4028
-keras_rs_nightly-0.3.1.dev202512130338.dist-info/METADATA,sha256=keexSmADAe0Wq2qD1bSC-g8JlUV17XrLU-zldDZ3ozM,5324
-keras_rs_nightly-0.3.1.dev202512130338.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.3.1.dev202512130338.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.3.1.dev202512130338.dist-info/RECORD,,
+keras_rs_nightly-0.4.1.dev202601250348.dist-info/METADATA,sha256=APNEvzS76AMD7Km5fXAWMnmc7VJspHKopOCOAd2xM1s,5324
+keras_rs_nightly-0.4.1.dev202601250348.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+keras_rs_nightly-0.4.1.dev202601250348.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.4.1.dev202601250348.dist-info/RECORD,,

{keras_rs_nightly-0.3.1.dev202512130338.dist-info → keras_rs_nightly-0.4.1.dev202601250348.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any