keras-rs-nightly 0.2.2.dev202507030337__tar.gz → 0.3.1.dev202511120334__tar.gz

{keras_rs_nightly-0.2.2.dev202507030337 → keras_rs_nightly-0.3.1.dev202511120334}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.2.2.dev202507030337
+Version: 0.3.1.dev202511120334
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0
@@ -8,8 +8,9 @@ Project-URL: Home, https://keras.io/keras_rs
 Project-URL: Repository, https://github.com/keras-team/keras-rs
 Classifier: Development Status :: 3 - Alpha
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Operating System :: Unix
 Classifier: Operating System :: Microsoft :: Windows
@@ -17,7 +18,7 @@ Classifier: Operating System :: MacOS
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: keras
 Requires-Dist: ml-dtypes

{keras_rs_nightly-0.2.2.dev202507030337 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/api/losses/__init__.py

@@ -4,6 +4,7 @@ This file was autogenerated. Do not edit it by hand,
 since your modifications would be overwritten.
 """
 
+from keras_rs.src.losses.list_mle_loss import ListMLELoss as ListMLELoss
 from keras_rs.src.losses.pairwise_hinge_loss import (
     PairwiseHingeLoss as PairwiseHingeLoss,
 )

{keras_rs_nightly-0.2.2.dev202507030337 → keras_rs_nightly-0.3.1.dev202511120334}/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(
@@ -146,14 +148,14 @@ class DistributedEmbedding(keras.layers.Layer):
     feature1 = keras_rs.layers.FeatureConfig(
         name="feature1",
         table=table1,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
     )
     feature2 = keras_rs.layers.FeatureConfig(
         name="feature2",
         table=table2,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
     )
 
     feature_configs = {
@@ -337,18 +339,33 @@ class DistributedEmbedding(keras.layers.Layer):
     embedding_layer = DistributedEmbedding(feature_configs)
 
     # Add preprocessing to a data input pipeline.
-    def train_dataset_generator():
-        for (inputs, weights), labels in iter(train_dataset):
+    def preprocessed_dataset_generator(dataset):
+        for (inputs, weights), labels in iter(dataset):
             yield embedding_layer.preprocess(
                 inputs, weights, training=True
             ), labels
 
-    preprocessed_train_dataset = train_dataset_generator()
+    preprocessed_train_dataset = preprocessed_dataset_generator(train_dataset)
     ```
     This explicit preprocessing stage combines the input and optional weights,
     so the new data can be passed directly into the `inputs` argument of the
     layer or model.
 
+    **NOTE**: When working in a multi-host setting with data parallelism, the
+    data needs to be sharded properly across hosts.  If the original dataset is
+    of type `tf.data.Dataset`, it will need to be manually sharded _prior_ to
+    applying the preprocess generator:
+    ```python
+    # Manually shard the dataset across hosts.
+    train_dataset = distribution.distribute_dataset(train_dataset)
+    distribution.auto_shard_dataset = False  # Dataset is already sharded.
+
+    # Add a preprocessing stage to the distributed data input pipeline.
+    train_dataset = preprocessed_dataset_generator(train_dataset)
+    ```
+    If the original dataset is _not_ a `tf.data.Dataset`, it must already be
+    pre-sharded across hosts.
+
     #### Usage in a Keras model
 
     Once the global distribution is set and the input preprocessing pipeline
@@ -503,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
@@ -575,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,
@@ -630,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]]]
@@ -669,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,
                 )
             )
@@ -679,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,
                 )
             )
@@ -765,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.
@@ -794,13 +822,13 @@ class DistributedEmbedding(keras.layers.Layer):
         table_stacking: str | Sequence[Sequence[str]],
     ) -> None:
         del table_stacking
-        table_to_embedding_layer: dict[TableConfig, EmbedReduce] = {}
+        table_config_id_to_embedding_layer: dict[int, EmbedReduce] = {}
         self._default_device_embedding_layers: dict[str, EmbedReduce] = {}
 
         for path, feature_config in feature_configs.items():
-            if feature_config.table in table_to_embedding_layer:
+            if id(feature_config.table) in table_config_id_to_embedding_layer:
                 self._default_device_embedding_layers[path] = (
-                    table_to_embedding_layer[feature_config.table]
+                    table_config_id_to_embedding_layer[id(feature_config.table)]
                 )
             else:
                 embedding_layer = EmbedReduce(
@@ -810,7 +838,9 @@ class DistributedEmbedding(keras.layers.Layer):
                     embeddings_initializer=feature_config.table.initializer,
                     combiner=feature_config.table.combiner,
                 )
-                table_to_embedding_layer[feature_config.table] = embedding_layer
+                table_config_id_to_embedding_layer[id(feature_config.table)] = (
+                    embedding_layer
+                )
                 self._default_device_embedding_layers[path] = embedding_layer
 
     def _default_device_build(
@@ -985,8 +1015,8 @@ class DistributedEmbedding(keras.layers.Layer):
 
         # The serialized `TableConfig` objects.
         table_config_dicts: list[dict[str, Any]] = []
-        # Mapping from `TableConfig` to index in `table_config_dicts`.
-        table_config_indices: dict[TableConfig, int] = {}
+        # Mapping from `TableConfig` id to index in `table_config_dicts`.
+        table_config_id_to_index: dict[int, int] = {}
 
         def serialize_feature_config(
             feature_config: FeatureConfig,
@@ -996,17 +1026,17 @@ class DistributedEmbedding(keras.layers.Layer):
             # key.
             feature_config_dict = feature_config.get_config()
 
-            if feature_config.table not in table_config_indices:
+            if id(feature_config.table) not in table_config_id_to_index:
                 # Save the serialized `TableConfig` the first time we see it and
                 # remember its index.
-                table_config_indices[feature_config.table] = len(
+                table_config_id_to_index[id(feature_config.table)] = len(
                     table_config_dicts
                 )
                 table_config_dicts.append(feature_config_dict["table"])
 
             # Replace the serialized `TableConfig` with its index.
-            feature_config_dict["table"] = table_config_indices[
-                feature_config.table
+            feature_config_dict["table"] = table_config_id_to_index[
+                id(feature_config.table)
             ]
             return feature_config_dict
 

{keras_rs_nightly-0.2.2.dev202507030337 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/src/layers/embedding/distributed_embedding_config.py

@@ -10,7 +10,7 @@ from keras_rs.src.api_export import keras_rs_export
 
 
 @keras_rs_export("keras_rs.layers.TableConfig")
-@dataclasses.dataclass(eq=True, unsafe_hash=True, order=True)
+@dataclasses.dataclass(order=True)
 class TableConfig:
     """Configuration for one embedding table.
 
@@ -88,7 +88,7 @@ class TableConfig:
 
 
 @keras_rs_export("keras_rs.layers.FeatureConfig")
-@dataclasses.dataclass(eq=True, unsafe_hash=True, order=True)
+@dataclasses.dataclass(order=True)
 class FeatureConfig:
     """Configuration for one embedding feature.
 
@@ -102,7 +102,10 @@ class FeatureConfig:
         input_shape: The input shape of the feature. The feature fed into the
             layer has to match the shape. Note that for ragged dimensions in the
             input, the dimension provided here presents the maximum value;
-            anything larger will be truncated.
+            anything larger will be truncated. Also note that the first
+            dimension represents the global batch size. For example, on TPU,
+            this represents the total number of samples that are dispatched to
+            all the TPUs connected to the current host.
         output_shape: The output shape of the feature activation. What is
             returned by the embedding layer has to match this shape.
     """