keras-rs-nightly 0.2.2.dev202508190331__py3-none-any.whl → 0.3.1.dev202512130338__py3-none-any.whl

keras_rs/src/losses/list_mle_loss.py

@@ -0,0 +1,212 @@
+from typing import Any
+
+import keras
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.metrics.utils import standardize_call_inputs_ranks
+
+
+@keras_rs_export("keras_rs.losses.ListMLELoss")
+class ListMLELoss(keras.losses.Loss):
+    """Implements ListMLE (Maximum Likelihood Estimation) loss for ranking.
+
+    ListMLE loss is a listwise ranking loss that maximizes the likelihood of
+    the ground truth ranking. It works by:
+    1. Sorting items by their relevance scores (labels)
+    2. Computing the probability of observing this ranking given the
+       predicted scores
+    3. Maximizing this likelihood (minimizing negative log-likelihood)
+
+    The loss is computed as the negative log-likelihood of the ground truth
+    ranking given the predicted scores:
+
+    ```
+    loss = -sum(log(exp(s_i) / sum(exp(s_j) for j >= i)))
+    ```
+
+    where s_i is the predicted score for item i in the sorted order.
+
+    Args:
+        temperature: Temperature parameter for scaling logits. Higher values
+            make the probability distribution more uniform. Defaults to 1.0.
+        reduction: Type of reduction to apply to the loss. In almost all cases
+            this should be `"sum_over_batch_size"`. Supported options are
+            `"sum"`, `"sum_over_batch_size"`, `"mean"`,
+            `"mean_with_sample_weight"` or `None`. Defaults to
+            `"sum_over_batch_size"`.
+        name: Optional name for the loss instance.
+        dtype: The dtype of the loss's computations. Defaults to `None`.
+
+    Examples:
+        ```python
+        # Basic usage
+        loss_fn = ListMLELoss()
+
+        # With temperature scaling
+        loss_fn = ListMLELoss(temperature=0.5)
+
+        # Example with synthetic data
+        y_true = [[3, 2, 1, 0]]  # Relevance scores
+        y_pred = [[0.8, 0.6, 0.4, 0.2]]  # Predicted scores
+        loss = loss_fn(y_true, y_pred)
+        ```
+    """
+
+    def __init__(self, temperature: float = 1.0, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+        if temperature <= 0.0:
+            raise ValueError(
+                f"`temperature` should be a positive float. Received: "
+                f"`temperature` = {temperature}."
+            )
+
+        self.temperature = temperature
+        self._epsilon = 1e-10
+
+    def compute_unreduced_loss(
+        self,
+        labels: types.Tensor,
+        logits: types.Tensor,
+        mask: types.Tensor | None = None,
+    ) -> tuple[types.Tensor, types.Tensor]:
+        """Compute the unreduced ListMLE loss.
+
+        Args:
+            labels: Ground truth relevance scores of
+                    shape [batch_size,list_size].
+            logits: Predicted scores of shape [batch_size, list_size].
+            mask: Optional mask of shape [batch_size, list_size].
+
+        Returns:
+            Tuple of (losses, weights) where losses has shape [batch_size, 1]
+            and weights has the same shape.
+        """
+
+        valid_mask = ops.greater_equal(labels, ops.cast(0.0, labels.dtype))
+
+        if mask is not None:
+            valid_mask = ops.logical_and(
+                valid_mask, ops.cast(mask, dtype="bool")
+            )
+
+        num_valid_items = ops.sum(
+            ops.cast(valid_mask, dtype=labels.dtype), axis=1, keepdims=True
+        )
+
+        batch_has_valid_items = ops.greater(num_valid_items, 0.0)
+
+        labels_for_sorting = ops.where(
+            valid_mask, labels, ops.full_like(labels, -1e9)
+        )
+        logits_masked = ops.where(
+            valid_mask, logits, ops.full_like(logits, -1e9)
+        )
+
+        sorted_logits, sorted_valid_mask = sort_by_scores(
+            tensors_to_sort=[logits_masked, valid_mask],
+            scores=labels_for_sorting,
+            mask=None,
+            shuffle_ties=False,
+            seed=None,
+        )
+        sorted_logits = ops.divide(
+            sorted_logits, ops.cast(self.temperature, dtype=sorted_logits.dtype)
+        )
+
+        valid_logits_for_max = ops.where(
+            sorted_valid_mask, sorted_logits, ops.full_like(sorted_logits, -1e9)
+        )
+        raw_max = ops.max(valid_logits_for_max, axis=1, keepdims=True)
+        raw_max = ops.where(
+            batch_has_valid_items, raw_max, ops.zeros_like(raw_max)
+        )
+        sorted_logits = ops.subtract(sorted_logits, raw_max)
+
+        # Set invalid positions to very negative BEFORE exp
+        sorted_logits = ops.where(
+            sorted_valid_mask, sorted_logits, ops.full_like(sorted_logits, -1e9)
+        )
+        exp_logits = ops.exp(sorted_logits)
+
+        reversed_exp = ops.flip(exp_logits, axis=1)
+        reversed_cumsum = ops.cumsum(reversed_exp, axis=1)
+        cumsum_from_right = ops.flip(reversed_cumsum, axis=1)
+
+        log_normalizers = ops.log(cumsum_from_right + self._epsilon)
+        log_probs = ops.subtract(sorted_logits, log_normalizers)
+
+        log_probs = ops.where(
+            sorted_valid_mask, log_probs, ops.zeros_like(log_probs)
+        )
+
+        negative_log_likelihood = ops.negative(
+            ops.sum(log_probs, axis=1, keepdims=True)
+        )
+
+        negative_log_likelihood = ops.where(
+            batch_has_valid_items,
+            negative_log_likelihood,
+            ops.zeros_like(negative_log_likelihood),
+        )
+
+        weights = ops.ones_like(negative_log_likelihood)
+
+        return negative_log_likelihood, weights
+
+    def call(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+    ) -> types.Tensor:
+        """Compute the ListMLE loss.
+
+        Args:
+            y_true: tensor or dict. Ground truth values. If tensor, of shape
+                `(list_size)` for unbatched inputs or `(batch_size, list_size)`
+                for batched inputs. If an item has a label of -1, it is ignored
+                in loss computation. If it is a dictionary, it should have two
+                keys: `"labels"` and `"mask"`. `"mask"` can be used to ignore
+                elements in loss computation.
+            y_pred: tensor. The predicted values, of shape `(list_size)` for
+                unbatched inputs or `(batch_size, list_size)` for batched
+                inputs. Should be of the same shape as `y_true`.
+
+        Returns:
+            The loss tensor of shape [batch_size].
+        """
+        mask = None
+        if isinstance(y_true, dict):
+            if "labels" not in y_true:
+                raise ValueError(
+                    '`"labels"` should be present in `y_true`. Received: '
+                    f"`y_true` = {y_true}"
+                )
+
+            mask = y_true.get("mask", None)
+            y_true = y_true["labels"]
+
+        y_true = ops.convert_to_tensor(y_true)
+        y_pred = ops.convert_to_tensor(y_pred)
+        if mask is not None:
+            mask = ops.convert_to_tensor(mask)
+
+        y_true, y_pred, mask, _ = standardize_call_inputs_ranks(
+            y_true, y_pred, mask
+        )
+
+        losses, weights = self.compute_unreduced_loss(
+            labels=y_true, logits=y_pred, mask=mask
+        )
+        losses = ops.multiply(losses, weights)
+        losses = ops.squeeze(losses, axis=-1)
+        return losses
+
+    # getting config
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update({"temperature": self.temperature})
+        return config

keras_rs/src/metrics/ranking_metrics_utils.py

@@ -85,6 +85,25 @@ def sort_by_scores(
     else:
         k = ops.minimum(k, max_possible_k)
 
+    # --- Work around for PyTorch instability ---
+    # Torch's `topk` is not stable with `sorted=True`, unlike JAX and TF.
+    # See:
+    #   - https://github.com/pytorch/pytorch/issues/27542
+    #   - https://github.com/pytorch/pytorch/issues/88227
+    #
+    # This small "stable offset" ensures deterministic tie-breaking for
+    # equal scores. We can remove this workaround once PyTorch adds a
+    # `stable=True` flag for topk.
+
+    if keras.backend.backend() == "torch" and not shuffle_ties:
+        list_size = ops.shape(scores)[1]
+        indices = ops.arange(list_size)
+        indices = ops.expand_dims(indices, axis=0)
+        indices = ops.broadcast_to(indices, ops.shape(scores))
+        stable_offset = ops.cast(indices, scores.dtype) * 1e-6
+        scores = ops.subtract(scores, stable_offset)
+    # --- End FIX ---
+
     # Shuffle ties randomly, and push masked values to the beginning.
     shuffled_indices = None
     if shuffle_ties or mask is not None:
@@ -205,12 +224,12 @@ def get_list_weights(
     return final_weights
 
 
-@keras.saving.register_keras_serializable()  # type: ignore[misc]
+@keras.saving.register_keras_serializable()  # type: ignore[untyped-decorator]
 def default_gain_fn(label: types.Tensor) -> types.Tensor:
     return ops.subtract(ops.power(2.0, label), 1.0)
 
 
-@keras.saving.register_keras_serializable()  # type: ignore[misc]
+@keras.saving.register_keras_serializable()  # type: ignore[untyped-decorator]
 def default_rank_discount_fn(rank: types.Tensor) -> types.Tensor:
     return ops.divide(
         ops.cast(1, dtype=rank.dtype),

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

{keras_rs_nightly-0.2.2.dev202508190331.dist-info → keras_rs_nightly-0.3.1.dev202512130338.dist-info}/METADATA

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

@@ -1,26 +1,26 @@
 keras_rs/__init__.py,sha256=8sjHiPN2GhUqAq4V7Vh4FLLqYw20-jgdI26ZKX5sg6M,350
 keras_rs/layers/__init__.py,sha256=ERqFu1R8FgeES5rO5QwauArbNCm8auj-AiCURtsG6Ro,1332
-keras_rs/losses/__init__.py,sha256=m04QOgxIUfJ2MvCUKLgEof-UbSNKgUYLPnY-D9NAclI,573
+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=q5U75eRIxARCwMEt9g7jww-oxkcHy2K7Yfe8YSBuaZs,224
+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=iJC02gxreAthInO2NrKbemRpz7-2OwCg1yAKSDeOWb0,45461
+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=nmDqZlXiZGxc2UmEaUULaOvRjwq427ReSltDGEBrv5I,5816
+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=jLqEuh_7hHM4bacyMY6FYpzlCUd9HcJHmnOG3ouucuQ,35734
-keras_rs/src/layers/embedding/jax/embedding_lookup.py,sha256=8LigXjPr7uQaUOdZM6yoLGoPYdRcbkXkFeL_sJoQ6uQ,8223
-keras_rs/src/layers/embedding/jax/embedding_utils.py,sha256=EHrQjPLl94STLWf9g8Ew8nuwupXRq-a_QmvFlXV6G6A,20331
+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=6aU7B-SX-WUqH8UC1amXg6BpFdalFmbr7rYQiiH_k4A,12883
-keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py,sha256=-K6l3PpIqOEVCzcVZN_cJMx9ddJnRKBTj486CP3wkps,17278
+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
@@ -31,6 +31,7 @@ keras_rs/src/layers/retrieval/remove_accidental_hits.py,sha256=WKoIhUSc6SvbgLXcS
 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
@@ -44,13 +45,14 @@ keras_rs/src/metrics/mean_reciprocal_rank.py,sha256=vr3ZZjpGYy2N-N7stcIm5elfHe9A
 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=0b03wiO9SjaHthtUYO4qezBFB8yLhFSwIRJhsL2fAJg,8785
+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_nightly-0.2.2.dev202508190331.dist-info/METADATA,sha256=v548uRNjECxcKkGmwpxB0_sTkeE9ERzVslyh9_zNb0Q,5273
-keras_rs_nightly-0.2.2.dev202508190331.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_rs_nightly-0.2.2.dev202508190331.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.2.2.dev202508190331.dist-info/RECORD,,
+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,,