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

keras_rs/src/losses/pairwise_mean_squared_error.py

@@ -0,0 +1,133 @@
+from keras import ops
+
+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.losses.pairwise_loss_utils import apply_pairwise_op
+
+
+@keras_rs_export("keras_rs.losses.PairwiseMeanSquaredError")
+class PairwiseMeanSquaredError(PairwiseLoss):
+    def pairwise_loss(self, pairwise_logits: types.Tensor) -> types.Tensor:
+        # Since we override `compute_unreduced_loss`, we do not need to
+        # implement this method.
+        pass
+
+    def compute_unreduced_loss(
+        self,
+        labels: types.Tensor,
+        logits: types.Tensor,
+        mask: types.Tensor | None = None,
+    ) -> tuple[types.Tensor, types.Tensor]:
+        # Override `PairwiseLoss.compute_unreduced_loss` since pairwise weights
+        # for MSE are computed differently.
+
+        batch_size, list_size = ops.shape(labels)
+
+        # Mask all values less than 0 (since less than 0 implies invalid
+        # labels).
+        valid_mask = ops.greater_equal(labels, ops.cast(0.0, labels.dtype))
+
+        if mask is not None:
+            valid_mask = ops.logical_and(valid_mask, mask)
+
+        # 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_diff = apply_pairwise_op(logits, ops.subtract)
+        valid_pair = apply_pairwise_op(valid_mask, ops.logical_and)
+        pairwise_mse = ops.square(pairwise_labels_diff - pairwise_logits_diff)
+
+        # Compute weights.
+        pairwise_weights = ops.ones_like(pairwise_mse)
+        # Exclude self pairs.
+        pairwise_weights = ops.subtract(
+            pairwise_weights,
+            ops.tile(ops.eye(list_size, list_size), (batch_size, 1, 1)),
+        )
+        # Include only valid pairs.
+        pairwise_weights = ops.multiply(
+            pairwise_weights, ops.cast(valid_pair, dtype=pairwise_weights.dtype)
+        )
+
+        return pairwise_mse, pairwise_weights
+
+
+formula = "loss = sum_{i} sum_{j} I(y_i > y_j) * (s_i - s_j)^2"
+explanation = """
+    - `(s_i - s_j)^2` is the squared difference between the predicted scores
+      of items `i` and `j`, which penalizes discrepancies between the predicted
+      order of items relative to their true order.
+"""
+extra_args = ""
+example = """
+    With `compile()` API:
+
+    ```python
+    model.compile(
+        loss=keras_rs.losses.PairwiseMeanSquaredError(),
+        ...
+    )
+    ```
+
+    As a standalone function with unbatched inputs:
+
+    >>> y_true = np.array([1.0, 0.0, 1.0, 3.0, 2.0])
+    >>> y_pred = np.array([1.0, 3.0, 2.0, 4.0, 0.8])
+    >>> pairwise_mse = keras_rs.losses.PairwiseMeanSquaredError()
+    >>> pairwise_mse(y_true=y_true, y_pred=y_pred)
+    >>> 19.10400
+
+    With batched inputs using default 'auto'/'sum_over_batch_size' reduction:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_mse = keras_rs.losses.PairwiseMeanSquaredError()
+    >>> pairwise_mse(y_true=y_true, y_pred=y_pred)
+    5.57999
+    
+    With masked inputs (useful for ragged inputs):
+
+    >>> y_true = {
+    ...     "labels": np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]]),
+    ...     "mask": np.array(
+    ...         [[True, True, True, True], [True, True, False, False]]
+    ...     ),
+    ... }
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_mse(y_true=y_true, y_pred=y_pred)
+    4.76000
+
+    With `sample_weight`:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> sample_weight = np.array(
+    ...     [[2.0, 3.0, 1.0, 1.0], [2.0, 1.0, 0.0, 0.0]]
+    ... )
+    >>> pairwise_mse = keras_rs.losses.PairwiseMeanSquaredError()
+    >>> pairwise_mse(
+    ...     y_true=y_true, y_pred=y_pred, sample_weight=sample_weight
+    ... )
+    11.0500
+
+    Using `'none'` reduction:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_mse = keras_rs.losses.PairwiseMeanSquaredError(
+    ...     reduction="none"
+    ... )
+    >>> pairwise_mse(y_true=y_true, y_pred=y_pred)
+    [[11., 17.,  5.,  5.], [2.04, 1.3199998, 1.6399999, 1.6399999]]
+"""
+
+PairwiseMeanSquaredError.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="mean squared error",
+    formula=formula,
+    explanation=explanation,
+    extra_args=extra_args,
+    example=example,
+)

keras_rs/src/losses/pairwise_soft_zero_one_loss.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.losses.pairwise_loss import PairwiseLoss
+from keras_rs.src.losses.pairwise_loss import pairwise_loss_subclass_doc_string
+
+
+@keras_rs_export("keras_rs.losses.PairwiseSoftZeroOneLoss")
+class PairwiseSoftZeroOneLoss(PairwiseLoss):
+    def pairwise_loss(self, pairwise_logits: types.Tensor) -> types.Tensor:
+        return ops.where(
+            ops.greater(pairwise_logits, ops.array(0.0)),
+            ops.subtract(ops.array(1.0), ops.sigmoid(pairwise_logits)),
+            ops.sigmoid(ops.negative(pairwise_logits)),
+        )
+
+
+formula = "loss = sum_{i} sum_{j} I(y_i > y_j) * (1 - sigmoid(s_i - s_j))"
+explanation = """
+    - `(1 - sigmoid(s_i - s_j))` represents the soft zero-one loss, which
+      approximates the ideal zero-one loss (which would be 1 if `s_i < s_j`
+      and 0 otherwise) with a smooth, differentiable function. This makes it
+      suitable for gradient-based optimization.
+"""
+extra_args = ""
+example = """
+    With `compile()` API:
+
+    ```python
+    model.compile(
+        loss=keras_rs.losses.PairwiseSoftZeroOneLoss(),
+        ...
+    )
+    ```
+
+    As a standalone function with unbatched inputs:
+
+    >>> y_true = np.array([1.0, 0.0, 1.0, 3.0, 2.0])
+    >>> y_pred = np.array([1.0, 3.0, 2.0, 4.0, 0.8])
+    >>> pairwise_soft_zero_one_loss = keras_rs.losses.PairwiseSoftZeroOneLoss()
+    >>> pairwise_soft_zero_one_loss(y_true=y_true, y_pred=y_pred)
+    0.86103
+
+    With batched inputs using default 'auto'/'sum_over_batch_size' reduction:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_soft_zero_one_loss = keras_rs.losses.PairwiseSoftZeroOneLoss()
+    >>> pairwise_soft_zero_one_loss(y_true=y_true, y_pred=y_pred)
+    0.46202
+
+    With masked inputs (useful for ragged inputs):
+
+    >>> y_true = {
+    ...     "labels": np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]]),
+    ...     "mask": np.array(
+    ...         [[True, True, True, True], [True, True, False, False]]
+    ...     ),
+    ... }
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_soft_zero_one_loss(y_true=y_true, y_pred=y_pred)
+    0.29468
+
+    With `sample_weight`:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> sample_weight = np.array(
+    ...     [[2.0, 3.0, 1.0, 1.0], [2.0, 1.0, 0.0, 0.0]]
+    ... )
+    >>> pairwise_soft_zero_one_loss = keras_rs.losses.PairwiseSoftZeroOneLoss()
+    >>> pairwise_soft_zero_one_loss(
+    ...     y_true=y_true, y_pred=y_pred, sample_weight=sample_weight
+    ... )
+    0.40478
+
+    Using `'none'` reduction:
+
+    >>> y_true = np.array([[1.0, 0.0, 1.0, 3.0], [0.0, 1.0, 2.0, 3.0]])
+    >>> y_pred = np.array([[1.0, 3.0, 2.0, 4.0], [1.0, 1.8, 2.0, 3.0]])
+    >>> pairwise_soft_zero_one_loss = keras_rs.losses.PairwiseSoftZeroOneLoss(
+    ...     reduction="none"
+    ... )
+    >>> pairwise_soft_zero_one_loss(y_true=y_true, y_pred=y_pred)
+    [
+        [0.8807971 , 0., 0.73105854, 0.43557024],
+        [0., 0.31002545, 0.7191075 , 0.61961967]
+    ]
+"""
+
+PairwiseSoftZeroOneLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="soft zero-one loss",
+    formula=formula,
+    explanation=explanation,
+    extra_args=extra_args,
+    example=example,
+)

keras_rs/src/metrics/dcg.py

@@ -0,0 +1,161 @@
+from typing import Any, Callable
+
+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_post_desc,
+)
+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: int | None = 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)`."""
+example = """
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 3, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> metric = keras_rs.metrics.DCG()(
+    ...     y_true=labels, y_pred=scores
+    ... )
+
+    Mask certain elements (can be used for uneven inputs):
+
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 3, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> mask = np.random.randint(0, 2, size=(batch_size, list_size), dtype=bool)
+    >>> metric = keras_rs.metrics.DCG()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+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_post_desc.format(
+    extra_args=extra_args, example=example
+)

keras_rs/src/metrics/mean_average_precision.py

@@ -0,0 +1,130 @@
+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_post_desc,
+)
+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=y_pred.dtype,
+        )
+        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 = ""
+example = """
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> metric = keras_rs.metrics.MeanAveragePrecision()(
+    ...     y_true=labels, y_pred=scores
+    ... )
+
+    Mask certain elements (can be used for uneven inputs):
+
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> mask = np.random.randint(0, 2, size=(batch_size, list_size), dtype=bool)
+    >>> metric = keras_rs.metrics.MeanAveragePrecision()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+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_post_desc.format(
+    extra_args=extra_args, example=example
+)

keras_rs/src/metrics/mean_reciprocal_rank.py

@@ -0,0 +1,121 @@
+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_post_desc,
+)
+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=y_pred.dtype,
+        )
+
+        # `reciprocal_rank = [1, 0.5, 0.33]`
+        reciprocal_rank = ops.divide(
+            ops.cast(1, dtype=y_pred.dtype),
+            ops.arange(1, list_length + 1, dtype=y_pred.dtype),
+        )
+
+        # `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=y_pred.dtype,
+        )
+        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 = ""
+example = """
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> metric = keras_rs.metrics.MeanReciprocalRank()(
+    ...     y_true=labels, y_pred=scores
+    ... )
+
+    Mask certain elements (can be used for uneven inputs):
+
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> mask = np.random.randint(0, 2, size=(batch_size, list_size), dtype=bool)
+    >>> metric = keras_rs.metrics.MeanReciprocalRank()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+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_post_desc.format(
+    extra_args=extra_args, example=example
+)