keras-rs-nightly 0.0.1.dev2025042903__tar.gz → 0.0.1.dev2025050103__tar.gz

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025042903
+Version: 0.0.1.dev2025050103
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0
@@ -93,7 +93,7 @@ Call `model.fit()` on dummy data:
 batch_size = 2
 x = np.random.randint(0, vocabulary_size, size=(batch_size,))
 y = np.random.random(size=(batch_size,))
-model.fit(input_data, y=y)
+model.fit(x, y=y)
 ```
 
 ### Use ranking losses and metrics

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/README.md

@@ -68,7 +68,7 @@ Call `model.fit()` on dummy data:
 batch_size = 2
 x = np.random.randint(0, vocabulary_size, size=(batch_size,))
 y = np.random.random(size=(batch_size,))
-model.fit(input_data, y=y)
+model.fit(x, y=y)
 ```
 
 ### Use ranking losses and metrics

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/layers/retrieval/hard_negative_mining.py

@@ -57,8 +57,9 @@ class HardNegativeMining(keras.layers.Layer):
         Returns:
             A tuple containing two tensors with the last dimension of
             `num_candidates` replaced with `num_hard_negatives + 1`.
-            - logits: `[..., num_hard_negatives + 1]` tensor of logits.
-            - labels: `[..., num_hard_negatives + 1]` one-hot tensor of labels.
+
+        * logits: `[..., num_hard_negatives + 1]` tensor of logits.
+        * labels: `[..., num_hard_negatives + 1]` one-hot tensor of labels.
         """
 
         # Number of sampled logits, i.e, the number of hard negatives to be

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/layers/retrieval/remove_accidental_hits.py

@@ -51,7 +51,7 @@ class RemoveAccidentalHits(keras.layers.Layer):
                 `labels`.
 
         Returns:
-            logits: The modified logits with the same shape as the input logits.
+            The modified logits with the same shape as the input logits.
         """
         # A more principled way is to implement
         # `softmax_cross_entropy_with_logits` with a input mask. Here we

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/layers/retrieval/sampling_probability_correction.py

@@ -50,7 +50,8 @@ class SamplingProbabilityCorrection(keras.layers.Layer):
             candidate_sampling_probability: The sampling probability with the
                 same shape as `logits`.
 
-        Returns: The corrected logits with the same shape as the input logits.
+        Returns:
+            The corrected logits with the same shape as the input logits.
         """
         return logits - ops.log(
             ops.clip(candidate_sampling_probability, self.epsilon, 1.0)

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/losses/pairwise_hinge_loss.py

@@ -14,13 +14,13 @@ class PairwiseHingeLoss(PairwiseLoss):
 
 formula = "loss = sum_{i} sum_{j} I(y_i > y_j) * max(0, 1 - (s_i - s_j))"
 explanation = """
-      - `max(0, 1 - (s_i - s_j))` is the hinge loss, which penalizes cases where
-        the score difference `s_i - s_j` is not sufficiently large when
-        `y_i > y_j`.
+    - `max(0, 1 - (s_i - s_j))` is the hinge loss, which penalizes cases where
+      the score difference `s_i - s_j` is not sufficiently large when
+      `y_i > y_j`.
 """
 extra_args = ""
 example = """
-    1. With `compile()` API:
+    With `compile()` API:
 
     ```python
     model.compile(
@@ -29,23 +29,24 @@ example = """
     )
     ```
 
-    2. As a standalone function:
-    2.1. Unbatched inputs
+    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_hinge_loss = keras_rs.losses.PairwiseHingeLoss()
     >>> pairwise_hinge_loss(y_true=y_true, y_pred=y_pred)
     2.32000
 
-    2.2 Batched inputs
-    2.2.1 Using default 'auto'/'sum_over_batch_size' reduction.
+    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_hinge_loss = keras_rs.losses.PairwiseHingeLoss()
     >>> pairwise_hinge_loss(y_true=y_true, y_pred=y_pred)
     0.75
 
-    2.2.2. With masked inputs (useful for ragged inputs)
+    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(
@@ -56,7 +57,8 @@ example = """
     >>> pairwise_hinge_loss(y_true=y_true, y_pred=y_pred)
     0.64999
 
-    2.2.3 With `sample_weight`
+    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(
@@ -68,7 +70,8 @@ example = """
     ... )
     1.02499
 
-    2.2.4 Using `'none'` reduction.
+    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_hinge_loss = keras_rs.losses.PairwiseHingeLoss(

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/losses/pairwise_logistic_loss.py

@@ -22,14 +22,14 @@ class PairwiseLogisticLoss(PairwiseLoss):
 
 formula = "loss = sum_{i} sum_{j} I(y_i > y_j) * log(1 + exp(-(s_i - s_j)))"
 explanation = """
-      - `log(1 + exp(-(s_i - s_j)))` is the logistic loss, which penalizes
-        cases where the score difference `s_i - s_j` is not sufficiently large
-        when `y_i > y_j`. This function provides a smooth approximation of the
-        ideal step function, making it suitable for gradient-based optimization.
+    - `log(1 + exp(-(s_i - s_j)))` is the logistic loss, which penalizes
+      cases where the score difference `s_i - s_j` is not sufficiently large
+      when `y_i > y_j`. This function provides a smooth approximation of the
+      ideal step function, making it suitable for gradient-based optimization.
 """
 extra_args = ""
 example = """
-    1. With `compile()` API:
+    With `compile()` API:
 
     ```python
     model.compile(
@@ -38,23 +38,24 @@ example = """
     )
     ```
 
-    2. As a standalone function:
-    2.1. Unbatched inputs
+    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_logistic_loss = keras_rs.losses.PairwiseLogisticLoss()
     >>> pairwise_logistic_loss(y_true=y_true, y_pred=y_pred)
     >>> 1.70708
 
-    2.2 Batched inputs
-    2.2.1 Using default 'auto'/'sum_over_batch_size' reduction.
+    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_logistic_loss = keras_rs.losses.PairwiseLogisticLoss()
     >>> pairwise_logistic_loss(y_true=y_true, y_pred=y_pred)
     0.73936
 
-    2.2.2. With masked inputs (useful for ragged inputs)
+    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(
@@ -65,7 +66,8 @@ example = """
     >>> pairwise_logistic_loss(y_true=y_true, y_pred=y_pred)
     0.53751
 
-    2.2.3 With `sample_weight`
+    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(
@@ -77,7 +79,8 @@ example = """
     ... )
     >>> 0.80337
 
-    2.2.4 Using `'none'` reduction.
+    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_logistic_loss = keras_rs.losses.PairwiseLogisticLoss(

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/losses/pairwise_loss.py

@@ -70,7 +70,8 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
         y_true: types.Tensor,
         y_pred: types.Tensor,
     ) -> types.Tensor:
-        """
+        """Compute the pairwise loss.
+
         Args:
             y_true: tensor or dict. Ground truth values. If tensor, of shape
                 `(list_size)` for unbatched inputs or `(batch_size, list_size)`
@@ -83,6 +84,9 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
             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.
         """
         mask = None
         if isinstance(y_true, dict):
@@ -134,11 +138,12 @@ pairwise_loss_subclass_doc_string = (
     ```
 
     where:
-      - `y_i` and `y_j` are the true labels of items `i` and `j`, respectively.
-      - `s_i` and `s_j` are the predicted scores of items `i` and `j`,
-        respectively.
-      - `I(y_i > y_j)` is an indicator function that equals 1 if `y_i > y_j`,
-        and 0 otherwise.{explanation}
+
+    - `y_i` and `y_j` are the true labels of items `i` and `j`, respectively.
+    - `s_i` and `s_j` are the predicted scores of items `i` and `j`,
+      respectively.
+    - `I(y_i > y_j)` is an indicator function that equals 1 if `y_i > y_j`,
+      and 0 otherwise.{explanation}
     Args:{extra_args}
         reduction: Type of reduction to apply to the loss. In almost all cases
             this should be `"sum_over_batch_size"`. Supported options are

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/losses/pairwise_mean_squared_error.py

@@ -59,13 +59,13 @@ class PairwiseMeanSquaredError(PairwiseLoss):
 
 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.
+    - `(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 = """
-    1. With `compile()` API:
+    With `compile()` API:
 
     ```python
     model.compile(
@@ -74,23 +74,24 @@ example = """
     )
     ```
 
-    2. As a standalone function:
-    2.1. Unbatched inputs
+    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
 
-    2.2 Batched inputs
-    2.2.1 Using default 'auto'/'sum_over_batch_size' reduction.
+    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
     
-    2.2.2. With masked inputs (useful for ragged inputs)
+    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(
@@ -101,7 +102,8 @@ example = """
     >>> pairwise_mse(y_true=y_true, y_pred=y_pred)
     4.76000
 
-    2.2.3 With `sample_weight`
+    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(
@@ -113,7 +115,8 @@ example = """
     ... )
     11.0500
 
-    2.2.4 Using `'none'` reduction.
+    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(

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/losses/pairwise_soft_zero_one_loss.py

@@ -18,14 +18,14 @@ class PairwiseSoftZeroOneLoss(PairwiseLoss):
 
 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.
+    - `(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 = """
-    1. With `compile()` API:
+    With `compile()` API:
 
     ```python
     model.compile(
@@ -34,23 +34,24 @@ example = """
     )
     ```
 
-    2. As a standalone function:
-    2.1. Unbatched inputs
+    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
 
-    2.2 Batched inputs
-    2.2.1 Using default 'auto'/'sum_over_batch_size' reduction.
+    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
 
-    2.2.2. With masked inputs (useful for ragged inputs)
+    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(
@@ -61,7 +62,8 @@ example = """
     >>> pairwise_soft_zero_one_loss(y_true=y_true, y_pred=y_pred)
     0.29468
 
-    2.2.3 With `sample_weight`
+    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(
@@ -73,7 +75,8 @@ example = """
     ... )
     0.40478
 
-    2.2.4 Using `'none'` reduction.
+    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(

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/dcg.py

@@ -111,16 +111,15 @@ 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.
-"""
+- `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`.
@@ -136,7 +135,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/mean_average_precision.py

@@ -25,7 +25,7 @@ class MeanAveragePrecision(RankingMetric):
     ) -> types.Tensor:
         relevance = ops.cast(
             ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         sorted_relevance, sorted_weights = sort_by_scores(
             tensors_to_sort=[relevance, sample_weight],
@@ -82,22 +82,18 @@ 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`.
-"""
+- `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
@@ -108,7 +104,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/mean_reciprocal_rank.py

@@ -44,13 +44,13 @@ class MeanReciprocalRank(RankingMetric):
             ops.greater_equal(
                 sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
             ),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
 
         # `reciprocal_rank = [1, 0.5, 0.33]`
         reciprocal_rank = ops.divide(
-            ops.cast(1, dtype="float32"),
-            ops.arange(1, list_length + 1, dtype="float32"),
+            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)`.
@@ -64,7 +64,7 @@ class MeanReciprocalRank(RankingMetric):
         # Get weights.
         overall_relevance = ops.cast(
             ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         per_list_weights = get_list_weights(
             weights=sample_weight, relevance=overall_relevance
@@ -95,7 +95,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/ndcg.py

@@ -109,7 +109,7 @@ class NDCG(RankingMetric):
 
 concept_sentence = (
     "It normalizes the Discounted Cumulative Gain (DCG) with the Ideal "
-    "Discounted Cumulative Gain (IDCG) for each list."
+    "Discounted Cumulative Gain (IDCG) for each list"
 )
 relevance_type = (
     "graded relevance scores (non-negative numbers where higher values "
@@ -124,11 +124,6 @@ score_range_interpretation = (
 )
 
 formula = """
-The metric calculates a weighted average nDCG score per list.
-For a single list, nDCG is computed as the ratio of the Discounted
-Cumulative Gain (DCG) of the predicted ranking to the Ideal Discounted
-Cumulative Gain (IDCG) of the best possible ranking:
-
 ```
 nDCG@k = DCG@k / IDCG@k
 ```
@@ -147,29 +142,25 @@ IDCG@k(y'') = sum_{i=1}^{k} (gain_fn(y''_i) / rank_discount_fn(i))
 ```
 
 where:
-    - `y'_i`: True relevance of the item at rank `i` in
-        the ranking induced by `y_pred`.
-    - `y''_i` True relevance of the item at rank `i` in
-        the *ideal* ranking (sorted by `y_true` descending).
-    - `gain_fn` is the user-provided function mapping relevance to gain.
-        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`
-        (1-based) to a discount value. The default function
-        (`default_rank_discount_fn`) is typically equivalent to
-        `lambda rank: 1 / log2(rank + 1)`.
-    - If IDCG@k is 0 (e.g., no relevant items), nDCG@k is defined as 0.
-    - The final result often aggregates these per-list nDCG scores,
-        potentially involving normalization by list-specific weights, to
-        produce a weighted average.
-"""
+- `y'_i`: True relevance of the item at rank `i` in the ranking induced by
+  `y_pred`.
+- `y''_i` True relevance of the item at rank `i` in the *ideal* ranking (sorted
+  by `y_true` descending).
+- `gain_fn` is the user-provided function mapping relevance to gain. 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` (1-based) to
+  a discount value. The default function (`default_rank_discount_fn`) is
+  typically equivalent to `lambda rank: 1 / log2(rank + 1)`.
+- If IDCG@k is 0 (e.g., no relevant items), nDCG@k is defined as 0.
+- The final result often aggregates these per-list nDCG scores, potentially
+  involving normalization by list-specific weights, to produce a weighted
+  average."""
 extra_args = """
-gain_fn: callable. Maps relevance scores (`y_true`) to gain values. The
-    default implements `2**y - 1`. Used for both DCG and IDCG.
-rank_discount_fn: callable. Maps rank positions (1-based) to discount
-    values. The default (`default_rank_discount_fn`) typically implements
-    `1 / log2(rank + 1)`. Used for both DCG and IDCG.
-"""
+        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
@@ -179,7 +170,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/precision_at_k.py

@@ -40,7 +40,7 @@ class PrecisionAtK(RankingMetric):
             ops.greater_equal(
                 sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
             ),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         list_length = ops.shape(sorted_y_true)[1]
         # TODO: We do not do this for MRR, and the other metrics. Do we need to
@@ -52,13 +52,13 @@ class PrecisionAtK(RankingMetric):
 
         per_list_precision = ops.divide_no_nan(
             ops.sum(relevance, axis=1, keepdims=True),
-            ops.cast(valid_list_length, dtype="float32"),
+            ops.cast(valid_list_length, dtype=y_pred.dtype),
         )
 
         # Get weights.
         overall_relevance = ops.cast(
             ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         per_list_weights = get_list_weights(
             weights=sample_weight, relevance=overall_relevance
@@ -91,7 +91,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/ranking_metric.py

@@ -116,6 +116,12 @@ class RankingMetric(keras.metrics.Mean, abc.ABC):
         if passed_mask is not None:
             passed_mask = ops.convert_to_tensor(passed_mask)
 
+        # Cast to the correct dtype.
+        y_true = ops.cast(y_true, dtype=self.dtype)
+        y_pred = ops.cast(y_pred, dtype=self.dtype)
+        if sample_weight is not None:
+            sample_weight = ops.cast(sample_weight, dtype=self.dtype)
+
         # === Process `sample_weight` ===
         if sample_weight is None:
             sample_weight = ops.cast(1, dtype=y_pred.dtype)
@@ -152,7 +158,7 @@ class RankingMetric(keras.metrics.Mean, abc.ABC):
 
         # Mask all values less than 0 (since less than 0 implies invalid
         # labels).
-        valid_mask = ops.greater_equal(y_true, ops.cast(0.0, y_true.dtype))
+        valid_mask = ops.greater_equal(y_true, ops.cast(0, y_true.dtype))
         if passed_mask is not None:
             valid_mask = ops.logical_and(valid_mask, passed_mask)
 
@@ -221,7 +227,6 @@ by sorting in descending order. {score_range_interpretation}.
 For each list of predicted scores `s` in `y_pred` and the corresponding list
 of true labels `y` in `y_true`, the per-query {metric_abbreviation} score is
 calculated as follows:
-
 {formula}
 
 The final {metric_abbreviation} score reported is typically the weighted

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/ranking_metrics_utils.py

@@ -163,7 +163,7 @@ def get_list_weights(
     # Identify lists where both weights and relevance sums are non-zero.
     nonzero_relevance = ops.cast(
         ops.logical_and(nonzero_weights, nonzero_relevance_condition),
-        dtype="float32",
+        dtype=weights.dtype,
     )
     # Count the number of lists with non-zero relevance and non-zero weights.
     nonzero_relevance_count = ops.sum(nonzero_relevance, axis=0, keepdims=True)
@@ -227,7 +227,7 @@ def compute_dcg(
     ] = default_rank_discount_fn,
 ) -> types.Tensor:
     list_size = ops.shape(y_true)[1]
-    positions = ops.arange(1, list_size + 1, dtype="float32")
+    positions = ops.arange(1, list_size + 1, dtype=y_true.dtype)
     gain = gain_fn(y_true)
     discount = rank_discount_fn(positions)
 

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/metrics/recall_at_k.py

@@ -38,11 +38,11 @@ class RecallAtK(RankingMetric):
             ops.greater_equal(
                 sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
             ),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         overall_relevance = ops.cast(
             ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
-            dtype="float32",
+            dtype=y_pred.dtype,
         )
         per_list_recall = ops.divide_no_nan(
             ops.sum(relevance, axis=1, keepdims=True),
@@ -82,7 +82,8 @@ example = """
     ...     y_true=labels, y_pred=scores
     ... )
 
-    # Mask certain elements (can be used for uneven inputs)
+    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))

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs/src/utils/doc_string_utils.py

@@ -30,8 +30,13 @@ def format_docstring(template: str, width: int = 80, **kwargs: Any) -> str:
                 textwrap.indent(formula_dedented, base_indent_str)
             )
         elif "where:" in stripped_block:
+            # Expect this to be already indented.
+            splitted_block = stripped_block.split("\n")
             processed_output.append(
-                textwrap.indent(stripped_block, base_indent_str)
+                textwrap.indent(
+                    splitted_block[0] + "\n\n" + "\n".join(splitted_block[1:]),
+                    base_indent_str,
+                )
             )
         else:
             processed_output.append(

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/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.dev2025042903"
+__version__ = "0.0.1.dev2025050103"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.0.1.dev2025042903 → keras_rs_nightly-0.0.1.dev2025050103}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025042903
+Version: 0.0.1.dev2025050103
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0
@@ -93,7 +93,7 @@ Call `model.fit()` on dummy data:
 batch_size = 2
 x = np.random.randint(0, vocabulary_size, size=(batch_size,))
 y = np.random.random(size=(batch_size,))
-model.fit(input_data, y=y)
+model.fit(x, y=y)
 ```
 
 ### Use ranking losses and metrics