keras-rs-nightly 0.3.1.dev202510280332__tar.gz → 0.3.1.dev202510300334__tar.gz

{keras_rs_nightly-0.3.1.dev202510280332 → keras_rs_nightly-0.3.1.dev202510300334}/PKG-INFO

@@ -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 → keras_rs_nightly-0.3.1.dev202510300334}/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_nightly-0.3.1.dev202510280332 → keras_rs_nightly-0.3.1.dev202510300334}/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 → keras_rs_nightly-0.3.1.dev202510300334}/keras_rs_nightly.egg-info/PKG-INFO

@@ -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