keras-rs-nightly 0.0.1.dev2025021903__py3-none-any.whl → 0.3.1.dev202512130338__py3-none-any.whl

keras_rs/src/utils/tpu_test_utils.py

@@ -0,0 +1,120 @@
+import contextlib
+import os
+import threading
+from typing import Any, Callable, ContextManager, Optional, Tuple, Union
+
+import keras
+import tensorflow as tf
+
+
+class DummyStrategy:
+    def scope(self) -> ContextManager[None]:
+        return contextlib.nullcontext()
+
+    @property
+    def num_replicas_in_sync(self) -> int:
+        return 1
+
+    def run(self, fn: Callable[..., Any], args: Tuple[Any, ...]) -> Any:
+        return fn(*args)
+
+    def experimental_distribute_dataset(
+        self, dataset: Any, options: Optional[Any] = None
+    ) -> Any:
+        del options
+        return dataset
+
+
+class JaxDummyStrategy(DummyStrategy):
+    @property
+    def num_replicas_in_sync(self) -> Any:
+        import jax
+
+        return jax.device_count("tpu")
+
+
+StrategyType = Union[tf.distribute.Strategy, DummyStrategy]
+
+_shared_strategy: Optional[StrategyType] = None
+_lock = threading.Lock()
+
+
+def create_tpu_strategy() -> Optional[StrategyType]:
+    """Initializes the TPU system and returns a TPUStrategy."""
+    print("Attempting to create TPUStrategy...")
+    try:
+        resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="")
+        tf.config.experimental_connect_to_cluster(resolver)
+        tf.tpu.experimental.initialize_tpu_system(resolver)
+        strategy = tf.distribute.TPUStrategy(resolver)
+        print(
+            "TPUStrategy created successfully."
+            "Devices: {strategy.extended.num_replicas_in_sync}"
+        )
+        return strategy
+    except Exception as e:
+        print(f"Error creating TPUStrategy: {e}")
+        return None
+
+
+def get_shared_tpu_strategy() -> Optional[StrategyType]:
+    """
+    Returns a session-wide shared TPUStrategy instance.
+    Creates the instance on the first call.
+    Returns None if not in a TPU environment or if creation fails.
+    """
+    global _shared_strategy
+    if _shared_strategy is not None:
+        return _shared_strategy
+
+    with _lock:
+        if _shared_strategy is None:
+            if "TPU_NAME" not in os.environ:
+                _shared_strategy = DummyStrategy()
+                return _shared_strategy
+            if keras.backend.backend() == "tensorflow":
+                resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
+                tf.config.experimental_connect_to_cluster(resolver)
+                topology = tf.tpu.experimental.initialize_tpu_system(resolver)
+                tpu_metadata = resolver.get_tpu_system_metadata()
+                device_assignment = tf.tpu.experimental.DeviceAssignment.build(
+                    topology, num_replicas=tpu_metadata.num_hosts
+                )
+                _shared_strategy = tf.distribute.TPUStrategy(
+                    resolver, experimental_device_assignment=device_assignment
+                )
+                print("### num_replicas", _shared_strategy.num_replicas_in_sync)
+            elif keras.backend.backend() == "jax":
+                _shared_strategy = JaxDummyStrategy()
+                print("### num_replicas", _shared_strategy.num_replicas_in_sync)
+            else:
+                _shared_strategy = DummyStrategy()
+    return _shared_strategy
+
+
+def run_with_strategy(
+    strategy: Any,
+    fn: Callable[..., Any],
+    *args: Any,
+    jit_compile: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """
+    Final wrapper fix: Flattens allowed kwargs into positional args before
+    entering tf.function to guarantee a fixed graph signature.
+    """
+    if keras.backend.backend() == "tensorflow":
+        all_inputs = (args, kwargs)
+
+        @tf.function(jit_compile=jit_compile)  # type: ignore[untyped-decorator]
+        def tf_function_wrapper(input_tuple: Tuple[Any, Any]) -> Any:
+            core_args, core_kwargs = input_tuple
+            if core_kwargs:
+                return strategy.run(fn, args=core_args, kwargs=core_kwargs)
+            else:
+                return strategy.run(fn, args=core_args)
+
+        return tf_function_wrapper(all_inputs)
+    else:
+        assert not jit_compile
+        return fn(*args, **kwargs)

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

{keras_rs_nightly-0.0.1.dev2025021903.dist-info → keras_rs_nightly-0.3.1.dev202512130338.dist-info}/METADATA

@@ -1,28 +1,34 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025021903
+Version: 0.3.1.dev202512130338
 Summary: Multi-backend recommender systems with Keras 3.
-Author-email: Keras RS team <keras-rs@google.com>
+Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0
-Project-URL: Home, https://keras.io/
+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.9
-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
 Classifier: Operating System :: MacOS
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
-Requires-Python: >=3.9
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: keras
+Requires-Dist: ml-dtypes
 
 # Keras Recommenders
 
+<p align="center" width="100%">
+<img src="https://i.imgur.com/m1BX7Zd.png" width="434" height="157" alt="KerasRS">
+</p>
+
 Keras Recommenders is a library for building recommender systems on top of
 Keras 3. Keras Recommenders works natively with TensorFlow, JAX, or PyTorch. It
 provides a collection of building blocks which help with the full workflow of
@@ -34,6 +40,80 @@ This library is an extension of the core Keras API; all high-level modules
 receive that same level of polish as core Keras. If you are familiar with Keras,
 congratulations! You already understand most of Keras Recommenders.
 
+## Quick Links
+
+- [Home page](https://keras.io/keras_rs)
+- [Examples](https://keras.io/keras_rs/examples)
+- [API documentation](https://keras.io/keras_rs/api)
+
+## Quickstart
+
+### Train your own cross network
+
+Choose a backend:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch"!
+```
+
+Import KerasRS and other libraries:
+
+```python
+import keras
+import keras_rs
+import numpy as np
+```
+
+Define a simple model using the `FeatureCross` layer:
+
+```python
+vocabulary_size = 32
+embedding_dim = 6
+
+inputs = keras.Input(shape=(), name='indices', dtype="int32")
+x0 = keras.layers.Embedding(
+    input_dim=vocabulary_size,
+    output_dim=embedding_dim
+)(inputs)
+x1 = keras_rs.layers.FeatureCross()(x0, x0)
+x2 = keras_rs.layers.FeatureCross()(x0, x1)
+output = keras.layers.Dense(units=10)(x2)
+model = keras.Model(inputs, output)
+```
+
+Compile the model:
+
+```python
+model.compile(
+    loss=keras.losses.MeanSquaredError(),
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4)
+)
+```
+
+Call `model.fit()` on dummy data:
+
+```python
+batch_size = 2
+x = np.random.randint(0, vocabulary_size, size=(batch_size,))
+y = np.random.random(size=(batch_size,))
+model.fit(x, y=y)
+```
+
+### Use ranking losses and metrics
+
+If your task is to rank items in a list, you can make use of the ranking losses
+and metrics which KerasRS provides. Below, we use the pairwise hinge loss and
+track the nDCG metric:
+
+```python
+model.compile(
+    loss=keras_rs.losses.PairwiseHingeLoss(),
+    metrics=[keras_rs.metrics.NDCG()],
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4),
+)
+```
+
 ## Installation
 
 Keras Recommenders is available on PyPI as `keras-rs`:
@@ -94,7 +174,7 @@ Here is the BibTeX entry:
 ```bibtex
 @misc{kerasrecommenders2024,
   title={KerasRecommenders},
-  author={Hertschuh, Fabien and  Chollet, Fran\c{c}ois and others},
+  author={Hertschuh, Fabien and  Chollet, Fran\c{c}ois and Sharma, Abheesht and others},
   year={2024},
   howpublished={\url{https://github.com/keras-team/keras-rs}},
 }

keras_rs_nightly-0.3.1.dev202512130338.dist-info/RECORD

@@ -0,0 +1,58 @@
+keras_rs/__init__.py,sha256=8sjHiPN2GhUqAq4V7Vh4FLLqYw20-jgdI26ZKX5sg6M,350
+keras_rs/layers/__init__.py,sha256=ERqFu1R8FgeES5rO5QwauArbNCm8auj-AiCURtsG6Ro,1332
+keras_rs/losses/__init__.py,sha256=WyyrxhWrayt-Hm6gSmZ5CPZifbPx0egDIothGi0Dpjk,646
+keras_rs/metrics/__init__.py,sha256=Qxpf6OFooIL9TIn2l3WgOea3HFRG0hq02glPAxtMZ9c,580
+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=E7EoQd2LJhzxrYtMh6GLQgtbm2hzCwUNr9AcZw0MOgc,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
+keras_rs/src/layers/embedding/distributed_embedding.py,sha256=94jxUHoGK3Gs9yfV0KxFTuqPo7XFnhgCNlO2FEeiSgM,1072
+keras_rs/src/layers/embedding/distributed_embedding_config.py,sha256=L41x6W1xcXI-3m94nOh_OsHn6OYpoynakKJvNboJnvE,5762
+keras_rs/src/layers/embedding/embed_reduce.py,sha256=c-MnEw1-KWs0jTf0JJ_ZBOY-9hRkiFyu989Dof3DnS8,12343
+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=lY7cwJgVUDpYUH-n7AcA1qlnSzhtckNtD3UqtpAYXz8,32267
+keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=a90tWTbU9tkFdESG3xir9PTtcvb1cmYR8vl5dDK9PSY,8703
+keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=5rQGli4Qflg0BU2-j_-4xbBxSqopqbtjkY2KKYWq64Y,7329
+keras_rs/src/layers/embedding/tensorflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+keras_rs/src/layers/embedding/tensorflow/config_conversion.py,sha256=HpuDthRQQ3X0EO8dW6OAdcgTODkujZlx_swgreVwXyk,13220
+keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=TBPYV8gP3ZnAFEwtxmWr_Rp3s-Cj0RrKzF6UOALJ4B0,17817
+keras_rs/src/layers/feature_interaction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+keras_rs/src/layers/feature_interaction/dot_interaction.py,sha256=Rs8xIHXNWQNiwjp_xzvQRmTSV1AyhJjDgVc3K5pTmrQ,8530
+keras_rs/src/layers/feature_interaction/feature_cross.py,sha256=Wq_eQvO0WTRlep69QbKi8TwY8bnFoF9vreP_j6ZHNFE,8666
+keras_rs/src/layers/retrieval/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=WVH9nHj6mGcK5DAOzYBGfUztPtGtMX4RKLac6GV1vLI,5644
+keras_rs/src/layers/retrieval/hard_negative_mining.py,sha256=n5UftRcuuR7Lh75vOdFdqatpsYqJDHCsraNtAjeWvoM,3575
+keras_rs/src/layers/retrieval/remove_accidental_hits.py,sha256=WKoIhUSc6SvbgLXcSqNvFUnkuyXfxWwsC7nAgYbON_U,3773
+keras_rs/src/layers/retrieval/retrieval.py,sha256=SFxMdooUhZy854SLZbpoyZR1Md4bHnpf7P077oVjjtU,4162
+keras_rs/src/layers/retrieval/sampling_probability_correction.py,sha256=3zD6LInxhyIvyujMleGqiuoPKsna2oaTN6JU6xMnW_M,1977
+keras_rs/src/losses/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+keras_rs/src/losses/list_mle_loss.py,sha256=NKRV_ZJUXFI1qG9_ugqxRyafRHreuUIokA7mbnsogBo,7433
+keras_rs/src/losses/pairwise_hinge_loss.py,sha256=tONOJpcwCw1mybwvyx8dAy5t6dDmlIn00enzWfQLXpQ,3049
+keras_rs/src/losses/pairwise_logistic_loss.py,sha256=40PFdCHDM7CLunT_PE3RbgxROVImw13dgVL3o3nzeNg,3473
+keras_rs/src/losses/pairwise_loss.py,sha256=Oydk8e7AGU0Mc9yvm6ccr_XDDfUe8EZlS4JJgyxKUm4,6197
+keras_rs/src/losses/pairwise_loss_utils.py,sha256=xvdGvdKNkvGvIaWYEQziWTFNa5EJz7rdkVGgrsnDHUk,1246
+keras_rs/src/losses/pairwise_mean_squared_error.py,sha256=sW-S3fUFtLV_cf8Tw9Wluzoe9t7lmXzSIAj6MXOpL2M,4823
+keras_rs/src/losses/pairwise_soft_zero_one_loss.py,sha256=YddVtJS8tKEeb0YrqGzEsr-6IDxH4uRjFrYkZDMWpkk,3492
+keras_rs/src/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+keras_rs/src/metrics/dcg.py,sha256=rUtWe8VwZ8unXIhPu7fjTiErpK6-IOtcaiuikjtY1VE,5823
+keras_rs/src/metrics/mean_average_precision.py,sha256=yUub4jGnqwTmxf694Z2ymRjMG_vO2HdyqqvDbcEhdSQ,4632
+keras_rs/src/metrics/mean_reciprocal_rank.py,sha256=vr3ZZjpGYy2N-N7stcIm5elfHe9A4W8uY4HADP8icMw,4046
+keras_rs/src/metrics/ndcg.py,sha256=ZBaKqV57K7jlto6ZVMxFNNRLdhzbLhdAR8TgDexjSjg,6922
+keras_rs/src/metrics/precision_at_k.py,sha256=Dj5R-rT_Yd5hAsk4f-BlNMujfgIdPXnFVGOw9u7BIZQ,4038
+keras_rs/src/metrics/ranking_metric.py,sha256=Lcl-Tt6HlI0f2wQpvAJ2M4mm5qCTZm-IgnLjjSEeNXE,10655
+keras_rs/src/metrics/ranking_metrics_utils.py,sha256=rtCy_8T5CbQ4UKr3ykFRaM_wZrQ-IqEJa_VnVMTH8nQ,9644
+keras_rs/src/metrics/recall_at_k.py,sha256=ssnQJC42KLN28cGrmzM-qR4M4iPqiQzWM2MfwYMq4ZE,3701
+keras_rs/src/metrics/utils.py,sha256=fGTo8j0ykVE5Y3yQCS2orSFcHY20Uxt0NazyPsybUsw,2471
+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/src/utils/tpu_test_utils.py,sha256=mQVBrI-CCBbXwQxBq1yDKGMwYhm4g4k3_AaSy_sCs0U,4028
+keras_rs_nightly-0.3.1.dev202512130338.dist-info/METADATA,sha256=keexSmADAe0Wq2qD1bSC-g8JlUV17XrLU-zldDZ3ozM,5324
+keras_rs_nightly-0.3.1.dev202512130338.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_rs_nightly-0.3.1.dev202512130338.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.3.1.dev202512130338.dist-info/RECORD,,

{keras_rs_nightly-0.0.1.dev2025021903.dist-info → keras_rs_nightly-0.3.1.dev202512130338.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

keras_rs/api/__init__.py

@@ -1,9 +0,0 @@
-"""DO NOT EDIT.
-
-This file was autogenerated. Do not edit it by hand,
-since your modifications would be overwritten.
-"""
-
-from keras_rs.api import layers
-from keras_rs.src.version import __version__
-from keras_rs.src.version import version

keras_rs/api/layers/__init__.py

@@ -1,11 +0,0 @@
-"""DO NOT EDIT.
-
-This file was autogenerated. Do not edit it by hand,
-since your modifications would be overwritten.
-"""
-
-from keras_rs.src.layers.modeling.dot_interaction import DotInteraction
-from keras_rs.src.layers.modeling.feature_cross import FeatureCross
-from keras_rs.src.layers.retrieval.brute_force_retrieval import (
-    BruteForceRetrieval,
-)

keras_rs_nightly-0.0.1.dev2025021903.dist-info/RECORD

@@ -1,19 +0,0 @@
-keras_rs/__init__.py,sha256=X3VNKb_6VDEs5GqcbEc_l8mAsefWb5UgSu8krnQdFcM,794
-keras_rs/api/__init__.py,sha256=BU34SgAwrjZh49ppXGxJPxreTOYW0C7vy3_x8nmvHUk,240
-keras_rs/api/layers/__init__.py,sha256=moxzBqnCCQsZiIgEUXn9mdtSALFTU5e0c4Ct8c0bWhg,355
-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=UyOdgjqrqg_b58opnY8n6gTiDHKVR8z_bmEruehERBk,514
-keras_rs/src/version.py,sha256=E-6oQfar5RwMDP5djal_81SRlf7inSYWDXk2dTmmAu4,222
-keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/modeling/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/modeling/dot_interaction.py,sha256=kzxs9IUI-llDfgqy-JUr-zBez9pA7A3zIG-O0lGxwWY,5097
-keras_rs/src/layers/modeling/feature_cross.py,sha256=nd2SBIe1gXC-vL_BeF3vrq8XqTAb6iLxhQGmtyUl3EM,8105
-keras_rs/src/layers/retrieval/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=mohILOt6PC6jHBztaowDbj3QBnSetuvkq55FmE39PlY,7321
-keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/utils/keras_utils.py,sha256=TUEpKHWDDBmzXJ0g0cTWYwQNjsHkPQ5HMvzNshLDOSU,897
-keras_rs_nightly-0.0.1.dev2025021903.dist-info/METADATA,sha256=MurKLV3gYnjxT2TNV9R3wYuRPfzlda5783DuBLd1S2E,3501
-keras_rs_nightly-0.0.1.dev2025021903.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-keras_rs_nightly-0.0.1.dev2025021903.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.0.1.dev2025021903.dist-info/RECORD,,