keras-rs-nightly 0.3.1.dev202509280329__tar.gz → 0.3.1.dev202509300325__tar.gz

{keras_rs_nightly-0.3.1.dev202509280329 → keras_rs_nightly-0.3.1.dev202509300325}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202509280329
+Version: 0.3.1.dev202509300325
 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.dev202509280329 → keras_rs_nightly-0.3.1.dev202509300325}/keras_rs/src/layers/embedding/jax/distributed_embedding.py

@@ -265,7 +265,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
         table_specs: Sequence[embedding_spec.TableSpec],
         num_shards: int,
         add_slot_variables: bool,
-    ) -> tuple[keras.Variable, tuple[keras.Variable, ...] | None]:
+    ) -> embedding.EmbeddingVariables:
         stacked_table_spec = typing.cast(
             embedding_spec.StackedTableSpec, table_specs[0].stacked_table_spec
         )
@@ -334,7 +334,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
                 slot_initializers, slot_variables
             )
 
-        return table_variable, slot_variables
+        return embedding.EmbeddingVariables(table_variable, slot_variables)
 
     @keras_utils.no_automatic_dependency_tracking
     def _sparsecore_init(
@@ -738,8 +738,8 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
         # Assign stacked table variables to the device values.
         keras.tree.map_structure_up_to(
             device_tables,
-            lambda table_and_slot_variables,
-            table_value: table_and_slot_variables[0].assign(table_value),
+            lambda embedding_variables,
+            table_value: embedding_variables.table.assign(table_value),
             self._table_and_slot_variables,
             device_tables,
         )
@@ -754,8 +754,10 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
 
         # Extract only the table variables, not the gradient slot variables.
         table_variables = {
-            name: jax.device_get(table_and_slots[0].value)
-            for name, table_and_slots in self._table_and_slot_variables.items()
+            name: jax.device_get(embedding_variables.table.value)
+            for name, embedding_variables in (
+                self._table_and_slot_variables.items()
+            )
         }
 
         return typing.cast(

{keras_rs_nightly-0.3.1.dev202509280329 → keras_rs_nightly-0.3.1.dev202509300325}/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.dev202509280329"
+__version__ = "0.3.1.dev202509300325"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.3.1.dev202509280329 → keras_rs_nightly-0.3.1.dev202509300325}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202509280329
+Version: 0.3.1.dev202509300325
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0