keras-rs-nightly 0.3.1.dev202510280332__py3-none-any.whl → 0.3.1.dev202510300334__py3-none-any.whl

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

@@ -1,5 +1,6 @@
 """JAX implementation of the TPU embedding layer."""
 
+import dataclasses
 import math
 import typing
 from typing import Any, Mapping, Sequence, Union
@@ -445,6 +446,30 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
         table_specs = embedding.get_table_specs(feature_specs)
         table_stacks = jte_table_stacking.get_table_stacks(table_specs)
 
+        # Create new instances of StackTableSpec with updated values that are
+        # the maximum from stacked tables.
+        stacked_table_specs = embedding.get_stacked_table_specs(feature_specs)
+        stacked_table_specs = {
+            stack_name: dataclasses.replace(
+                stacked_table_spec,
+                max_ids_per_partition=max(
+                    table.max_ids_per_partition
+                    for table in table_stacks[stack_name]
+                ),
+                max_unique_ids_per_partition=max(
+                    table.max_unique_ids_per_partition
+                    for table in table_stacks[stack_name]
+                ),
+            )
+            for stack_name, stacked_table_spec in stacked_table_specs.items()
+        }
+
+        # Rewrite the stacked_table_spec in all TableSpecs.
+        for stack_name, table_specs in table_stacks.items():
+            stacked_table_spec = stacked_table_specs[stack_name]
+            for table_spec in table_specs:
+                table_spec.stacked_table_spec = stacked_table_spec
+
         # Create variables for all stacked tables and slot variables.
         with sparsecore_distribution.scope():
             self._table_and_slot_variables = {

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.dev202510280332"
+__version__ = "0.3.1.dev202510300334"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.3.1.dev202510280332.dist-info → keras_rs_nightly-0.3.1.dev202510300334.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202510280332
+Version: 0.3.1.dev202510300334
 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.dev202510280332.dist-info → keras_rs_nightly-0.3.1.dev202510300334.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=LBNXhlFa6P1nQhY9SUb1spImnuIwMrfVQVmJfnUfGGM,224
+keras_rs/src/version.py,sha256=TSXfr6OYMhKNF7BAmoXC-DV8ckEnzJyki8qYNvkM-Rk,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
@@ -15,7 +15,7 @@ keras_rs/src/layers/embedding/embed_reduce.py,sha256=c-MnEw1-KWs0jTf0JJ_ZBOY-9hR
 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=O3G0AFRzukYdXPRyx7ZDqDvNgJrcbFwTCYTHigfdiKw,29628
+keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=d0QEb_0JHIs_cqYvmkowfpC1EWuMGsG-nwqggjPiSYY,30710
 keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=8LigXjPr7uQaUOdZM6yoLGoPYdRcbkXkFeL_sJoQ6uQ,8223
 keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=slJ0XwkI1z4vTAnRXQwm39LFnK9AL3CODuGRn5BufgE,8292
 keras_rs/src/layers/embedding/tensorflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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.3.1.dev202510280332.dist-info/METADATA,sha256=VVWxJMaJj1ItnuQmyfTAFAAdcBJrQ-sxUAg1SYs6c8Q,5324
-keras_rs_nightly-0.3.1.dev202510280332.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.3.1.dev202510280332.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.3.1.dev202510280332.dist-info/RECORD,,
+keras_rs_nightly-0.3.1.dev202510300334.dist-info/METADATA,sha256=6rVE87pwIQxWpXVZ6oNm659wv9MgUJLzX-bpi4mrS3o,5324
+keras_rs_nightly-0.3.1.dev202510300334.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.3.1.dev202510300334.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.3.1.dev202510300334.dist-info/RECORD,,