keras-rs-nightly 0.2.2.dev202508140338__py3-none-any.whl → 0.2.2.dev202508150339__py3-none-any.whl

keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -1,4 +1,5 @@
 import collections
+import dataclasses
 import importlib.util
 import typing
 from typing import Any, Sequence
@@ -20,9 +21,10 @@ EmbedReduce = embed_reduce.EmbedReduce
 SUPPORTED_PLACEMENTS = ("auto", "default_device", "sparsecore")
 
 
-PlacementAndPath = collections.namedtuple(
-    "PlacementAndPath", ["placement", "path"]
-)
+@dataclasses.dataclass(eq=True, unsafe_hash=True, order=True)
+class PlacementAndPath:
+    placement: str
+    path: str
 
 
 def _ragged_to_dense_inputs(
@@ -518,12 +520,12 @@ class DistributedEmbedding(keras.layers.Layer):
         With these structures in place, the steps to:
         - go from the deeply nested structure to the two-level structure are:
           - `assert_same_struct` as `self._feature_configs`
-          - `flatten`
-          - `pack_sequence_as` `self._placement_to_path_to_feature_config`
+          - use `self._feature_deeply_nested_placement_and_paths` to map from
+            deeply nested to two-level
         - go from the two-level structure to the deeply nested structure:
-         - `assert_same_struct` as `self._placement_to_path_to_feature_config`
-         - `flatten`
-         - `pack_sequence_as` `self._feature_configs`
+          - `assert_same_struct` as `self._placement_to_path_to_feature_config`
+          - use `self._feature_deeply_nested_placement_and_paths` to locate each
+            output in the two-level dicts
 
         Args:
             feature_configs: The deeply nested structure of `FeatureConfig` or
@@ -590,14 +592,14 @@ class DistributedEmbedding(keras.layers.Layer):
         ] = collections.defaultdict(dict)
 
         def populate_placement_to_path_to_input_shape(
-            placement_and_path: PlacementAndPath, input_shape: types.Shape
+            pp: PlacementAndPath, input_shape: types.Shape
         ) -> None:
-            placement_to_path_to_input_shape[placement_and_path.placement][
-                placement_and_path.path
-            ] = input_shape
+            placement_to_path_to_input_shape[pp.placement][pp.path] = (
+                input_shape
+            )
 
         keras.tree.map_structure_up_to(
-            self._feature_configs,
+            self._feature_deeply_nested_placement_and_paths,
             populate_placement_to_path_to_input_shape,
             self._feature_deeply_nested_placement_and_paths,
             input_shapes,
@@ -645,35 +647,40 @@ class DistributedEmbedding(keras.layers.Layer):
         """
         # Verify input structure.
         keras.tree.assert_same_structure(self._feature_configs, inputs)
+        if weights is not None:
+            keras.tree.assert_same_structure(self._feature_configs, weights)
 
         if not self.built:
-            input_shapes = keras.tree.map_structure_up_to(
-                self._feature_configs,
+            input_shapes = keras.tree.map_structure(
                 lambda array: backend.standardize_shape(array.shape),
                 inputs,
             )
             self.build(input_shapes)
 
-        # Go from deeply nested structure of inputs to flat inputs.
-        flat_inputs = keras.tree.flatten(inputs)
+        # Go from deeply nested to nested dict placement -> path -> input.
+        def to_placement_to_path(
+            tensors: types.Nested[types.Tensor],
+        ) -> dict[str, dict[str, types.Tensor]]:
+            result: dict[str, dict[str, types.Tensor]] = {
+                p: dict() for p in self._placement_to_path_to_feature_config
+            }
 
-        # Go from flat to nested dict placement -> path -> input.
-        placement_to_path_to_inputs = keras.tree.pack_sequence_as(
-            self._placement_to_path_to_feature_config, flat_inputs
-        )
+            def populate(pp: PlacementAndPath, x: types.Tensor) -> None:
+                result[pp.placement][pp.path] = x
 
-        if weights is not None:
-            # Same for weights if present.
-            keras.tree.assert_same_structure(self._feature_configs, weights)
-            flat_weights = keras.tree.flatten(weights)
-            placement_to_path_to_weights = keras.tree.pack_sequence_as(
-                self._placement_to_path_to_feature_config, flat_weights
+            keras.tree.map_structure(
+                populate,
+                self._feature_deeply_nested_placement_and_paths,
+                tensors,
             )
-        else:
-            # Populate keys for weights.
-            placement_to_path_to_weights = {
-                k: None for k in placement_to_path_to_inputs
-            }
+            return result
+
+        placement_to_path_to_inputs = to_placement_to_path(inputs)
+
+        # Same for weights if present.
+        placement_to_path_to_weights = (
+            to_placement_to_path(weights) if weights is not None else None
+        )
 
         placement_to_path_to_preprocessed: dict[
             str, dict[str, dict[str, types.Nested[types.Tensor]]]
@@ -684,7 +691,9 @@ class DistributedEmbedding(keras.layers.Layer):
             placement_to_path_to_preprocessed["sparsecore"] = (
                 self._sparsecore_preprocess(
                     placement_to_path_to_inputs["sparsecore"],
-                    placement_to_path_to_weights["sparsecore"],
+                    placement_to_path_to_weights["sparsecore"]
+                    if placement_to_path_to_weights is not None
+                    else None,
                     training,
                 )
             )
@@ -694,7 +703,9 @@ class DistributedEmbedding(keras.layers.Layer):
             placement_to_path_to_preprocessed["default_device"] = (
                 self._default_device_preprocess(
                     placement_to_path_to_inputs["default_device"],
-                    placement_to_path_to_weights["default_device"],
+                    placement_to_path_to_weights["default_device"]
+                    if placement_to_path_to_weights is not None
+                    else None,
                     training,
                 )
             )
@@ -780,11 +791,13 @@ class DistributedEmbedding(keras.layers.Layer):
             placement_to_path_to_outputs,
         )
 
-        # Go from placement -> path -> output to flat outputs.
-        flat_outputs = keras.tree.flatten(placement_to_path_to_outputs)
+        # Go from placement -> path -> output to deeply nested structure.
+        def populate_output(pp: PlacementAndPath) -> types.Tensor:
+            return placement_to_path_to_outputs[pp.placement][pp.path]
 
-        # Go from flat outputs to deeply nested structure.
-        return keras.tree.pack_sequence_as(self._feature_configs, flat_outputs)
+        return keras.tree.map_structure(
+            populate_output, self._feature_deeply_nested_placement_and_paths
+        )
 
     def get_embedding_tables(self) -> dict[str, types.Tensor]:
         """Return the content of the embedding tables by table name.

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.dev202508140338"
+__version__ = "0.2.2.dev202508150339"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.2.2.dev202508140338.dist-info → keras_rs_nightly-0.2.2.dev202508150339.dist-info}/METADATA

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

@@ -5,10 +5,10 @@ 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=ippx-4zzdkzAo2PfQ781mExG7U_MuF_NU2uSV_czqaQ,224
+keras_rs/src/version.py,sha256=wvBc0HOPE8xIqLsp6cWz-5TtgtVINmg5amOyutxMFiY,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=J1mW7sgOgfkR8s-n3uyuqMb_iwPgrD6j0J5OBSfXQt8,45006
+keras_rs/src/layers/embedding/base_distributed_embedding.py,sha256=iJC02gxreAthInO2NrKbemRpz7-2OwCg1yAKSDeOWb0,45461
 keras_rs/src/layers/embedding/distributed_embedding.py,sha256=94jxUHoGK3Gs9yfV0KxFTuqPo7XFnhgCNlO2FEeiSgM,1072
 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
@@ -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.dev202508140338.dist-info/METADATA,sha256=0fe9-5D5lrJlcYw0mutzFG7GcwD_I_lZPce5Q7SKL_g,5273
-keras_rs_nightly-0.2.2.dev202508140338.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.2.2.dev202508140338.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.2.2.dev202508140338.dist-info/RECORD,,
+keras_rs_nightly-0.2.2.dev202508150339.dist-info/METADATA,sha256=TqXJCEpnATn3zSL0bi-7-MadI5ojXvgUpQN4EzkhKkA,5273
+keras_rs_nightly-0.2.2.dev202508150339.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.2.2.dev202508150339.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.2.2.dev202508150339.dist-info/RECORD,,