keras-rs-nightly 0.0.1.dev2025042103__tar.gz → 0.0.1.dev2025042503__tar.gz

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025042103
+Version: 0.0.1.dev2025042503
 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
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 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

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/api/__init__.py

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

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/api/layers/__init__.py

@@ -5,19 +5,21 @@ since your modifications would be overwritten.
 """
 
 from keras_rs.src.layers.feature_interaction.dot_interaction import (
-    DotInteraction,
+    DotInteraction as DotInteraction,
+)
+from keras_rs.src.layers.feature_interaction.feature_cross import (
+    FeatureCross as FeatureCross,
 )
-from keras_rs.src.layers.feature_interaction.feature_cross import FeatureCross
 from keras_rs.src.layers.retrieval.brute_force_retrieval import (
-    BruteForceRetrieval,
+    BruteForceRetrieval as BruteForceRetrieval,
 )
 from keras_rs.src.layers.retrieval.hard_negative_mining import (
-    HardNegativeMining,
+    HardNegativeMining as HardNegativeMining,
 )
 from keras_rs.src.layers.retrieval.remove_accidental_hits import (
-    RemoveAccidentalHits,
+    RemoveAccidentalHits as RemoveAccidentalHits,
 )
-from keras_rs.src.layers.retrieval.retrieval import Retrieval
+from keras_rs.src.layers.retrieval.retrieval import Retrieval as Retrieval
 from keras_rs.src.layers.retrieval.sampling_probability_correction import (
-    SamplingProbabilityCorrection,
+    SamplingProbabilityCorrection as SamplingProbabilityCorrection,
 )

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/api/losses/__init__.py

@@ -0,0 +1,18 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras_rs.src.losses.pairwise_hinge_loss import (
+    PairwiseHingeLoss as PairwiseHingeLoss,
+)
+from keras_rs.src.losses.pairwise_logistic_loss import (
+    PairwiseLogisticLoss as PairwiseLogisticLoss,
+)
+from keras_rs.src.losses.pairwise_mean_squared_error import (
+    PairwiseMeanSquaredError as PairwiseMeanSquaredError,
+)
+from keras_rs.src.losses.pairwise_soft_zero_one_loss import (
+    PairwiseSoftZeroOneLoss as PairwiseSoftZeroOneLoss,
+)

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/api/metrics/__init__.py

@@ -0,0 +1,16 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras_rs.src.metrics.dcg import DCG as DCG
+from keras_rs.src.metrics.mean_average_precision import (
+    MeanAveragePrecision as MeanAveragePrecision,
+)
+from keras_rs.src.metrics.mean_reciprocal_rank import (
+    MeanReciprocalRank as MeanReciprocalRank,
+)
+from keras_rs.src.metrics.ndcg import NDCG as NDCG
+from keras_rs.src.metrics.precision_at_k import PrecisionAtK as PrecisionAtK
+from keras_rs.src.metrics.recall_at_k import RecallAtK as RecallAtK

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/src/losses/pairwise_hinge_loss.py

@@ -20,6 +20,7 @@ explanation = """
 """
 extra_args = ""
 PairwiseHingeLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="hinge loss",
     formula=formula,
     explanation=explanation,
     extra_args=extra_args,

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/src/losses/pairwise_logistic_loss.py

@@ -29,6 +29,7 @@ explanation = """
 """
 extra_args = ""
 PairwiseLogisticLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="logistic loss",
     formula=formula,
     explanation=explanation,
     extra_args=extra_args,

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/src/losses/pairwise_loss.py

@@ -1,12 +1,12 @@
 import abc
-from typing import Optional
+from typing import Any, Optional
 
 import keras
 from keras import ops
 
 from keras_rs.src import types
-from keras_rs.src.utils.pairwise_loss_utils import pairwise_comparison
-from keras_rs.src.utils.pairwise_loss_utils import process_loss_call_inputs
+from keras_rs.src.losses.pairwise_loss_utils import pairwise_comparison
+from keras_rs.src.metrics.utils import standardize_call_inputs_ranks
 
 
 class PairwiseLoss(keras.losses.Loss, abc.ABC):
@@ -22,14 +22,22 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
     `pairwise_loss` method.
     """
 
-    # TODO: Add `temperature`, `lambda_weights`.
+    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
+
+        # TODO(abheesht): Add `lambda_weights`.
 
     @abc.abstractmethod
     def pairwise_loss(self, pairwise_logits: types.Tensor) -> types.Tensor:
-        raise NotImplementedError(
-            "All subclasses of `keras_rs.losses.pairwise_loss.PairwiseLoss`"
-            "must implement the `pairwise_loss()` method."
-        )
+        pass
 
     def compute_unreduced_loss(
         self,
@@ -50,6 +58,10 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
             mask=valid_mask,
             logits_op=ops.subtract,
         )
+        pairwise_logits = ops.divide(
+            pairwise_logits,
+            ops.cast(self.temperature, dtype=pairwise_logits.dtype),
+        )
 
         return self.pairwise_loss(pairwise_logits), pairwise_labels
 
@@ -66,8 +78,8 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
                 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, i.e., pairs will not be formed
-                with those items. Note that the final mask is an and of the
-                passed mask, and `labels == -1`.
+                with those items. Note that the final mask is an `and` of the
+                passed mask, and `labels >= 0`.
             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`.
@@ -83,7 +95,14 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
             mask = y_true.get("mask", None)
             y_true = y_true["labels"]
 
-        y_true, y_pred, mask = process_loss_call_inputs(y_true, y_pred, mask)
+        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
@@ -92,9 +111,14 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
         losses = ops.sum(losses, axis=-1)
         return losses
 
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update({"temperature": self.temperature})
+        return config
+
 
 pairwise_loss_subclass_doc_string = (
-    "Computes pairwise hinge loss between true labels and predicted scores."
+    "Computes pairwise {loss_name} between true labels and predicted scores."
     """
     This loss function is designed for ranking tasks, where the goal is to
     correctly order items within each list. It computes the loss by comparing

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/src/losses/pairwise_loss_utils.py

@@ -0,0 +1,39 @@
+from typing import Callable
+
+from keras import ops
+
+from keras_rs.src import types
+
+
+def apply_pairwise_op(
+    x: types.Tensor, op: Callable[[types.Tensor, types.Tensor], types.Tensor]
+) -> types.Tensor:
+    return op(
+        ops.expand_dims(x, axis=-1),
+        ops.expand_dims(x, axis=-2),
+    )
+
+
+def pairwise_comparison(
+    labels: types.Tensor,
+    logits: types.Tensor,
+    mask: types.Tensor,
+    logits_op: Callable[[types.Tensor, types.Tensor], types.Tensor],
+) -> tuple[types.Tensor, types.Tensor]:
+    # Compute the difference for all pairs in a list. The output is a tensor
+    # with shape `(batch_size, list_size, list_size)`, where `[:, i, j]` stores
+    # information for pair `(i, j)`.
+    pairwise_labels_diff = apply_pairwise_op(labels, ops.subtract)
+    pairwise_logits = apply_pairwise_op(logits, logits_op)
+
+    # Keep only those cases where `l_i < l_j`.
+    pairwise_labels = ops.cast(
+        ops.greater(pairwise_labels_diff, 0), dtype=labels.dtype
+    )
+    if mask is not None:
+        valid_pairs = apply_pairwise_op(mask, ops.logical_and)
+        pairwise_labels = ops.multiply(
+            pairwise_labels, ops.cast(valid_pairs, dtype=pairwise_labels.dtype)
+        )
+
+    return pairwise_labels, pairwise_logits

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/src/losses/pairwise_mean_squared_error.py

@@ -6,7 +6,7 @@ from keras_rs.src import types
 from keras_rs.src.api_export import keras_rs_export
 from keras_rs.src.losses.pairwise_loss import PairwiseLoss
 from keras_rs.src.losses.pairwise_loss import pairwise_loss_subclass_doc_string
-from keras_rs.src.utils.pairwise_loss_utils import apply_pairwise_op
+from keras_rs.src.losses.pairwise_loss_utils import apply_pairwise_op
 
 
 @keras_rs_export("keras_rs.losses.PairwiseMeanSquaredError")
@@ -65,6 +65,7 @@ explanation = """
 """
 extra_args = ""
 PairwiseMeanSquaredError.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="mean squared error",
     formula=formula,
     explanation=explanation,
     extra_args=extra_args,

{keras_rs_nightly-0.0.1.dev2025042103 → keras_rs_nightly-0.0.1.dev2025042503}/keras_rs/src/losses/pairwise_soft_zero_one_loss.py

@@ -25,6 +25,7 @@ explanation = """
 """
 extra_args = ""
 PairwiseSoftZeroOneLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="soft zero-one loss",
     formula=formula,
     explanation=explanation,
     extra_args=extra_args,

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/src/metrics/dcg.py

@@ -0,0 +1,140 @@
+from typing import Any, Callable, Optional
+
+from keras import ops
+from keras.saving import deserialize_keras_object
+from keras.saving import serialize_keras_object
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.metrics.ranking_metric import RankingMetric
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string,
+)
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string_args,
+)
+from keras_rs.src.metrics.ranking_metrics_utils import compute_dcg
+from keras_rs.src.metrics.ranking_metrics_utils import default_gain_fn
+from keras_rs.src.metrics.ranking_metrics_utils import default_rank_discount_fn
+from keras_rs.src.metrics.ranking_metrics_utils import get_list_weights
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.utils.doc_string_utils import format_docstring
+
+
+@keras_rs_export("keras_rs.metrics.DCG")
+class DCG(RankingMetric):
+    def __init__(
+        self,
+        k: Optional[int] = None,
+        gain_fn: Callable[[types.Tensor], types.Tensor] = default_gain_fn,
+        rank_discount_fn: Callable[
+            [types.Tensor], types.Tensor
+        ] = default_rank_discount_fn,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(k=k, **kwargs)
+
+        self.gain_fn = gain_fn
+        self.rank_discount_fn = rank_discount_fn
+
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        sorted_y_true, sorted_weights = sort_by_scores(
+            tensors_to_sort=[y_true, sample_weight],
+            scores=y_pred,
+            k=self.k,
+            mask=mask,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+
+        dcg = compute_dcg(
+            y_true=sorted_y_true,
+            sample_weight=sorted_weights,
+            gain_fn=self.gain_fn,
+            rank_discount_fn=self.rank_discount_fn,
+        )
+
+        per_list_weights = get_list_weights(
+            weights=sample_weight, relevance=self.gain_fn(y_true)
+        )
+        # Since we have already multiplied with `sample_weight`, we need to
+        # divide by `per_list_weights` so as to nullify the multiplication
+        # which `keras.metrics.Mean` will do.
+        per_list_dcg = ops.divide_no_nan(dcg, per_list_weights)
+
+        return per_list_dcg, per_list_weights
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update(
+            {
+                "gain_fn": serialize_keras_object(self.gain_fn),
+                "rank_discount_fn": serialize_keras_object(
+                    self.rank_discount_fn
+                ),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "DCG":
+        config["gain_fn"] = deserialize_keras_object(config["gain_fn"])
+        config["rank_discount_fn"] = deserialize_keras_object(
+            config["rank_discount_fn"]
+        )
+        return cls(**config)
+
+
+concept_sentence = (
+    "It computes the sum of the graded relevance scores of items, applying a "
+    "configurable discount based on position"
+)
+relevance_type = (
+    "graded relevance scores (non-negative numbers where higher values "
+    "indicate greater relevance)"
+)
+score_range_interpretation = (
+    "Scores are non-negative, with higher values indicating better ranking "
+    "quality (highly relevant items are ranked higher). The score for a single "
+    "list is not bounded or normalized, i.e., it does not lie in a range"
+)
+
+formula = """
+```
+DCG@k(y', w') = sum_{i=1}^{k} (gain_fn(y'_i) / rank_discount_fn(i))
+```
+
+where:
+    - `y'_i` is the true relevance score of the item ranked at position `i`
+        (obtained by sorting `y_true` according to `y_pred`).
+    - `gain_fn` is the user-provided function mapping relevance `y'_i` to a
+        gain value. The default function (`default_gain_fn`) is typically
+        equivalent to `lambda y: 2**y - 1`.
+    - `rank_discount_fn` is the user-provided function mapping rank `i`
+        to a discount value. The default function (`default_rank_discount_fn`)
+        is typically equivalent to `lambda rank: 1 / log2(rank + 1)`.
+    - The final result aggregates these per-list scores.
+"""
+extra_args = """
+        gain_fn: callable. Maps relevance scores (`y_true`) to gain values. The
+            default implements `2**y - 1`.
+        rank_discount_fn: function. Maps rank positions to discount
+            values. The default (`default_rank_discount_fn`) implements
+            `1 / log2(rank + 1)`."""
+
+DCG.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Discounted Cumulative Gain",
+    metric_abbreviation="DCG",
+    concept_sentence=concept_sentence,
+    relevance_type=relevance_type,
+    score_range_interpretation=score_range_interpretation,
+    formula=formula,
+) + ranking_metric_subclass_doc_string_args.format(extra_args=extra_args)

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/src/metrics/mean_average_precision.py

@@ -0,0 +1,112 @@
+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_metric import RankingMetric
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string,
+)
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string_args,
+)
+from keras_rs.src.metrics.ranking_metrics_utils import get_list_weights
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.utils.doc_string_utils import format_docstring
+
+
+@keras_rs_export("keras_rs.metrics.MeanAveragePrecision")
+class MeanAveragePrecision(RankingMetric):
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        relevance = ops.cast(
+            ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
+            dtype="float32",
+        )
+        sorted_relevance, sorted_weights = sort_by_scores(
+            tensors_to_sort=[relevance, sample_weight],
+            scores=y_pred,
+            mask=mask,
+            k=self.k,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+        per_list_relevant_counts = ops.cumsum(sorted_relevance, axis=1)
+        per_list_cutoffs = ops.cumsum(ops.ones_like(sorted_relevance), axis=1)
+        per_list_precisions = ops.divide_no_nan(
+            per_list_relevant_counts, per_list_cutoffs
+        )
+
+        total_precision = ops.sum(
+            ops.multiply(
+                per_list_precisions,
+                ops.multiply(sorted_weights, sorted_relevance),
+            ),
+            axis=1,
+            keepdims=True,
+        )
+
+        # Compute the total relevance.
+        total_relevance = ops.sum(
+            ops.multiply(sample_weight, relevance), axis=1, keepdims=True
+        )
+
+        per_list_map = ops.divide_no_nan(total_precision, total_relevance)
+
+        per_list_weights = get_list_weights(sample_weight, relevance)
+
+        return per_list_map, per_list_weights
+
+
+concept_sentence = (
+    "It calculates the average of precision values computed after each "
+    "relevant item present in the ranked list"
+)
+relevance_type = "binary indicators (0 or 1) of relevance"
+score_range_interpretation = (
+    "Scores range from 0 to 1, with higher values indicating that relevant "
+    "items are generally positioned higher in the ranking"
+)
+
+formula = """
+The formula for average precision is defined below. MAP is the mean over average
+precision computed for each list.
+
+```
+AP(y, s) = sum_j (P@j(y, s) * rel(j)) / sum_i y_i
+rel(j) = y_i if rank(s_i) = j
+```
+
+where:
+    - `j` represents the rank position (starting from 1).
+    - `sum_j` indicates a summation over all ranks `j` from 1 up to the list
+        size (or `k`).
+    - `P@j(y, s)` denotes the Precision at rank `j`, calculated as the
+        number of relevant items found within the top `j` positions divided by
+        `j`.
+    - `rel(j)` represents the relevance of the item specifically at rank
+        `j`. `rel(j)` is 1 if the item at rank `j` is relevant, and 0
+        otherwise.
+    - `y_i` is the true relevance label of the original item `i` before
+        ranking.
+    - `rank(s_i)` is the rank position assigned to item `i` based on its
+        score `s_i`.
+    - `sum_i y_i` calculates the total number of relevant items in the
+        original list `y`.
+"""
+extra_args = ""
+
+MeanAveragePrecision.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Mean Average Precision",
+    metric_abbreviation="MAP",
+    concept_sentence=concept_sentence,
+    relevance_type=relevance_type,
+    score_range_interpretation=score_range_interpretation,
+    formula=formula,
+) + ranking_metric_subclass_doc_string_args.format(extra_args=extra_args)

keras_rs_nightly-0.0.1.dev2025042503/keras_rs/src/metrics/mean_reciprocal_rank.py

@@ -0,0 +1,98 @@
+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_metric import RankingMetric
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string,
+)
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string_args,
+)
+from keras_rs.src.metrics.ranking_metrics_utils import get_list_weights
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.utils.doc_string_utils import format_docstring
+
+
+@keras_rs_export("keras_rs.metrics.MeanReciprocalRank")
+class MeanReciprocalRank(RankingMetric):
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        # Assume: `y_true = [0, 0, 1]`, `y_pred = [0.1, 0.9, 0.2]`.
+        # `sorted_y_true = [0, 1, 0]` (sorted in descending order).
+        (sorted_y_true,) = sort_by_scores(
+            tensors_to_sort=[y_true],
+            scores=y_pred,
+            mask=mask,
+            k=self.k,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+
+        # This will depend on `k`, i.e., it will not always be the same as
+        # `len(y_true)`.
+        list_length = ops.shape(sorted_y_true)[1]
+
+        # We consider only binary relevance here, anything above 1 is treated
+        # as 1. `relevance = [0., 1., 0.]`.
+        relevance = ops.cast(
+            ops.greater_equal(
+                sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
+            ),
+            dtype="float32",
+        )
+
+        # `reciprocal_rank = [1, 0.5, 0.33]`
+        reciprocal_rank = ops.divide(
+            ops.cast(1, dtype="float32"),
+            ops.arange(1, list_length + 1, dtype="float32"),
+        )
+
+        # `mrr` should be of shape `(batch_size, 1)`.
+        # `mrr = amax([0., 0.5, 0.]) = 0.5`
+        mrr = ops.amax(
+            ops.multiply(relevance, reciprocal_rank),
+            axis=1,
+            keepdims=True,
+        )
+
+        # Get weights.
+        overall_relevance = ops.cast(
+            ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
+            dtype="float32",
+        )
+        per_list_weights = get_list_weights(
+            weights=sample_weight, relevance=overall_relevance
+        )
+
+        return mrr, per_list_weights
+
+
+concept_sentence = (
+    "It focuses on the rank position of the single highest-scoring relevant "
+    "item"
+)
+relevance_type = "binary indicators (0 or 1) of relevance"
+score_range_interpretation = (
+    "Scores range from 0 to 1, with 1 indicating the first relevant item was "
+    "always ranked first"
+)
+formula = """```
+MRR(y, s) = max_{i} y_{i} / rank(s_{i})
+```"""
+extra_args = ""
+MeanReciprocalRank.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Mean Reciprocal Rank",
+    metric_abbreviation="MRR",
+    concept_sentence=concept_sentence,
+    relevance_type=relevance_type,
+    score_range_interpretation=score_range_interpretation,
+    formula=formula,
+) + ranking_metric_subclass_doc_string_args.format(extra_args=extra_args)