keras-rs-nightly 0.3.1.dev202510220333__py3-none-any.whl → 0.3.1.dev202510240328__py3-none-any.whl

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

@@ -9,6 +9,7 @@ import keras
 import numpy as np
 from jax import numpy as jnp
 from jax.experimental import layout as jax_layout
+from jax.experimental import multihost_utils
 from jax_tpu_embedding.sparsecore.lib.nn import embedding
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 from jax_tpu_embedding.sparsecore.lib.nn import (
@@ -600,31 +601,26 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             # underlying stacked tables specs in the feature specs.
 
             # Aggregate stats across all processes/devices via pmax.
-            num_local_cpu_devices = jax.local_device_count("cpu")
-
-            def pmax_aggregate(x: Any) -> Any:
-                if not hasattr(x, "ndim"):
-                    x = np.array(x)
-                tiled_x = np.tile(x, (num_local_cpu_devices, *([1] * x.ndim)))
-                return jax.pmap(
-                    lambda y: jax.lax.pmax(y, "all_cpus"),  # type: ignore[no-untyped-call]
-                    axis_name="all_cpus",
-                    backend="cpu",
-                )(tiled_x)[0]
-
-            full_stats = jax.tree.map(pmax_aggregate, stats)
+            all_stats = multihost_utils.process_allgather(stats)
+            aggregated_stats = jax.tree.map(
+                lambda x: jnp.max(x, axis=0), all_stats
+            )
 
             # Check if stats changed enough to warrant action.
             stacked_table_specs = embedding.get_stacked_table_specs(
                 self._config.feature_specs
             )
             changed = any(
-                np.max(full_stats.max_ids_per_partition[stack_name])
+                np.max(aggregated_stats.max_ids_per_partition[stack_name])
                 > spec.max_ids_per_partition
-                or np.max(full_stats.max_unique_ids_per_partition[stack_name])
+                or np.max(
+                    aggregated_stats.max_unique_ids_per_partition[stack_name]
+                )
                 > spec.max_unique_ids_per_partition
                 or (
-                    np.max(full_stats.required_buffer_size_per_sc[stack_name])
+                    np.max(
+                        aggregated_stats.required_buffer_size_per_sc[stack_name]
+                    )
                     * num_sc_per_device
                 )
                 > (spec.suggested_coo_buffer_size_per_device or 0)
@@ -634,7 +630,9 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             # Update configuration and repeat preprocessing if stats changed.
             if changed:
                 embedding.update_preprocessing_parameters(
-                    self._config.feature_specs, full_stats, num_sc_per_device
+                    self._config.feature_specs,
+                    aggregated_stats,
+                    num_sc_per_device,
                 )
 
                 # Re-execute preprocessing with consistent input statistics.

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

{keras_rs_nightly-0.3.1.dev202510220333.dist-info → keras_rs_nightly-0.3.1.dev202510240328.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202510220333
+Version: 0.3.1.dev202510240328
 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.dev202510220333.dist-info → keras_rs_nightly-0.3.1.dev202510240328.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=5K8isi168MtN0VSp5_c1VYmRMw8pLpswFkARLUS9Z8Q,224
+keras_rs/src/version.py,sha256=PT5g4Jbeo7CzD3xiVH7sJhFwaI8vEYnZJmnZMryylNo,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=l3gpFXBAdkJw7yVntl0s25excCfC5jryyqBxUKZd2Fk,29820
+keras_rs/src/layers/embedding/jax/distributed_embedding.py,sha256=O3G0AFRzukYdXPRyx7ZDqDvNgJrcbFwTCYTHigfdiKw,29628
 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.dev202510220333.dist-info/METADATA,sha256=xiV9gYIXS_a1sbBEFQwlfLn4mZ3GyAmfBFVMYRhC4Jc,5324
-keras_rs_nightly-0.3.1.dev202510220333.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.3.1.dev202510220333.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.3.1.dev202510220333.dist-info/RECORD,,
+keras_rs_nightly-0.3.1.dev202510240328.dist-info/METADATA,sha256=Ezi7EQGTudw0xaxAsoHu4thA66EKIeXFNOBvHTBkH1I,5324
+keras_rs_nightly-0.3.1.dev202510240328.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.3.1.dev202510240328.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.3.1.dev202510240328.dist-info/RECORD,,