keras-rs-nightly 0.0.1.dev2025042803__tar.gz → 0.0.1.dev2025043003__tar.gz

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025042803
+Version: 0.0.1.dev2025043003
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0
@@ -36,6 +36,80 @@ This library is an extension of the core Keras API; all high-level modules
 receive that same level of polish as core Keras. If you are familiar with Keras,
 congratulations! You already understand most of Keras Recommenders.
 
+## Quick Links
+
+- [Home page](https://keras.io/keras_rs)
+- [Examples](https://keras.io/keras_rs/examples)
+- [API documentation](https://keras.io/keras_rs/api)
+
+## Quickstart
+
+### Train your own cross network
+
+Choose a backend:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch"!
+```
+
+Import KerasRS and other libraries:
+
+```python
+import keras
+import keras_rs
+import numpy as np
+```
+
+Define a simple model using the `FeatureCross` layer:
+
+```python
+vocabulary_size = 32
+embedding_dim = 6
+
+inputs = keras.Input(shape=(), name='indices', dtype="int32")
+x0 = keras.layers.Embedding(
+    input_dim=vocabulary_size,
+    output_dim=embedding_dim
+)(inputs)
+x1 = keras_rs.layers.FeatureCross()(x0, x0)
+x2 = keras_rs.layers.FeatureCross()(x0, x1)
+output = keras.layers.Dense(units=10)(x2)
+model = keras.Model(inputs, output)
+```
+
+Compile the model:
+
+```python
+model.compile(
+    loss=keras.losses.MeanSquaredError(),
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4)
+)
+```
+
+Call `model.fit()` on dummy data:
+
+```python
+batch_size = 2
+x = np.random.randint(0, vocabulary_size, size=(batch_size,))
+y = np.random.random(size=(batch_size,))
+model.fit(x, y=y)
+```
+
+### Use ranking losses and metrics
+
+If your task is to rank items in a list, you can make use of the ranking losses
+and metrics which KerasRS provides. Below, we use the pairwise hinge loss and
+track the nDCG metric:
+
+```python
+model.compile(
+    loss=keras_rs.losses.PairwiseHingeLoss(),
+    metrics=[keras_rs.metrics.NDCG()],
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4),
+)
+```
+
 ## Installation
 
 Keras Recommenders is available on PyPI as `keras-rs`:

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/README.md

@@ -11,6 +11,80 @@ This library is an extension of the core Keras API; all high-level modules
 receive that same level of polish as core Keras. If you are familiar with Keras,
 congratulations! You already understand most of Keras Recommenders.
 
+## Quick Links
+
+- [Home page](https://keras.io/keras_rs)
+- [Examples](https://keras.io/keras_rs/examples)
+- [API documentation](https://keras.io/keras_rs/api)
+
+## Quickstart
+
+### Train your own cross network
+
+Choose a backend:
+
+```python
+import os
+os.environ["KERAS_BACKEND"] = "jax"  # Or "tensorflow" or "torch"!
+```
+
+Import KerasRS and other libraries:
+
+```python
+import keras
+import keras_rs
+import numpy as np
+```
+
+Define a simple model using the `FeatureCross` layer:
+
+```python
+vocabulary_size = 32
+embedding_dim = 6
+
+inputs = keras.Input(shape=(), name='indices', dtype="int32")
+x0 = keras.layers.Embedding(
+    input_dim=vocabulary_size,
+    output_dim=embedding_dim
+)(inputs)
+x1 = keras_rs.layers.FeatureCross()(x0, x0)
+x2 = keras_rs.layers.FeatureCross()(x0, x1)
+output = keras.layers.Dense(units=10)(x2)
+model = keras.Model(inputs, output)
+```
+
+Compile the model:
+
+```python
+model.compile(
+    loss=keras.losses.MeanSquaredError(),
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4)
+)
+```
+
+Call `model.fit()` on dummy data:
+
+```python
+batch_size = 2
+x = np.random.randint(0, vocabulary_size, size=(batch_size,))
+y = np.random.random(size=(batch_size,))
+model.fit(x, y=y)
+```
+
+### Use ranking losses and metrics
+
+If your task is to rank items in a list, you can make use of the ranking losses
+and metrics which KerasRS provides. Below, we use the pairwise hinge loss and
+track the nDCG metric:
+
+```python
+model.compile(
+    loss=keras_rs.losses.PairwiseHingeLoss(),
+    metrics=[keras_rs.metrics.NDCG()],
+    optimizer=keras.optimizers.Adam(learning_rate=3e-4),
+)
+```
+
 ## Installation
 
 Keras Recommenders is available on PyPI as `keras-rs`:

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/keras_rs/src/layers/feature_interaction/dot_interaction.py

@@ -30,6 +30,53 @@ class DotInteraction(keras.layers.Layer):
             but is much slower.
         **kwargs: Args to pass to the base class.
 
+    Example:
+
+    ```python
+    # 1. Simple forward pass
+    batch_size = 2
+    embedding_dim = 32
+    feature1 = np.random.randn(batch_size, embedding_dim)
+    feature2 = np.random.randn(batch_size, embedding_dim)
+    feature3 = np.random.randn(batch_size, embedding_dim)
+    feature_interactions = keras_rs.layers.DotInteraction()(
+        [feature1, feature2, feature3]
+    )
+
+    # 2. After embedding layer in a model
+    vocabulary_size = 32
+    embedding_dim = 6
+
+    # Create a simple model containing the layer.
+    feature_input_1 = keras.Input(shape=(), name='indices_1', dtype="int32")
+    feature_input_2 = keras.Input(shape=(), name='indices_2', dtype="int32")
+    feature_input_3 = keras.Input(shape=(), name='indices_3', dtype="int32")
+    x1 = keras.layers.Embedding(
+        input_dim=vocabulary_size,
+        output_dim=embedding_dim
+    )(feature_input_1)
+    x2 = keras.layers.Embedding(
+        input_dim=vocabulary_size,
+        output_dim=embedding_dim
+    )(feature_input_2)
+    x3 = keras.layers.Embedding(
+        input_dim=vocabulary_size,
+        output_dim=embedding_dim
+    )(feature_input_3)
+    feature_interactions = keras_rs.layers.DotInteraction()([x1, x2, x3])
+    output = keras.layers.Dense(units=10)(x2)
+    model = keras.Model(
+        [feature_input_1, feature_input_2, feature_input_3], output
+    )
+
+    # Call the model on the inputs.
+    batch_size = 2
+    f1 = np.random.randint(0, vocabulary_size, size=(batch_size,))
+    f2 = np.random.randint(0, vocabulary_size, size=(batch_size,))
+    f3 = np.random.randint(0, vocabulary_size, size=(batch_size,))
+    outputs = model([f1, f2, f3])
+    ```
+
     References:
     - [M. Naumov et al.](https://arxiv.org/abs/1906.00091)
     """

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/keras_rs/src/layers/feature_interaction/feature_cross.py

@@ -57,13 +57,32 @@ class FeatureCross(keras.layers.Layer):
     Example:
 
     ```python
-    # after embedding layer in a functional model
-    input = keras.Input(shape=(), name='indices', dtype="int64")
-    x0 = keras.layers.Embedding(input_dim=32, output_dim=6)(x0)
-    x1 = FeatureCross()(x0, x0)
-    x2 = FeatureCross()(x0, x1)
+    # 1. Simple forward pass
+    batch_size = 2
+    embedding_dim = 32
+    feature1 = np.random.randn(batch_size, embedding_dim)
+    feature2 = np.random.randn(batch_size, embedding_dim)
+    crossed_features = keras_rs.layers.FeatureCross()(feature1, feature2)
+
+    # 2. After embedding layer in a model
+    vocabulary_size = 32
+    embedding_dim = 6
+
+    # Create a simple model containing the layer.
+    inputs = keras.Input(shape=(), name='indices', dtype="int32")
+    x0 = keras.layers.Embedding(
+        input_dim=vocabulary_size,
+        output_dim=embedding_dim
+    )(inputs)
+    x1 = keras_rs.layers.FeatureCross()(x0, x0)
+    x2 = keras_rs.layers.FeatureCross()(x0, x1)
     logits = keras.layers.Dense(units=10)(x2)
-    model = keras.Model(input, logits)
+    model = keras.Model(inputs, logits)
+
+    # Call the model on the inputs.
+    batch_size = 2
+    input_data = np.random.randint(0, vocabulary_size, size=(batch_size,))
+    outputs = model(input_data)
     ```
 
     References:

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/keras_rs/src/layers/retrieval/hard_negative_mining.py

@@ -12,11 +12,27 @@ MAX_FLOAT = ml_dtypes.finfo("float32").max / 100.0
 
 @keras_rs_export("keras_rs.layers.HardNegativeMining")
 class HardNegativeMining(keras.layers.Layer):
-    """Transforms logits and labels to return hard negatives.
+    """Filter logits and labels to return hard negatives.
+
+    The output will include logits and labels for the requested number of hard
+    negatives as well as the positive candidate.
 
     Args:
         num_hard_negatives: How many hard negatives to return.
         **kwargs: Args to pass to the base class.
+
+    Example:
+
+    ```python
+    # Create layer with the configured number of hard negatives to mine.
+    hard_negative_mining = keras_rs.layers.HardNegativeMining(
+        num_hard_negatives=10
+    )
+
+    # This will retrieve the top 10 negative candidates plus the positive
+    # candidate from `labels` for each row.
+    out_logits, out_labels = hard_negative_mining(in_logits, in_labels)
+    ```
     """
 
     def __init__(self, num_hard_negatives: int, **kwargs: Any) -> None:
@@ -33,15 +49,17 @@ class HardNegativeMining(keras.layers.Layer):
         negatives as well as the positive candidate.
 
         Args:
-            logits: logits tensor, typically `[batch_size, num_candidates]` but
-                can have more dimensions or be 1D as `[num_candidates]`.
-            labels: one-hot labels tensor, must be the same shape as `logits`.
+            logits: The logits tensor, typically `[batch_size, num_candidates]`
+                but can have more dimensions or be 1D as `[num_candidates]`.
+            labels: The one-hot labels tensor, must be the same shape as
+                `logits`.
 
         Returns:
-            tuple containing two tensors with the last dimension of
+            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.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/keras_rs/src/layers/retrieval/remove_accidental_hits.py

@@ -15,6 +15,18 @@ class RemoveAccidentalHits(keras.layers.Layer):
 
     Zeroes the logits of negative candidates that have the same ID as the
     positive candidate in that row.
+
+    Example:
+
+    ```python
+    # Create layer with the configured number of hard negatives to mine.
+    remove_accidental_hits = keras_rs.layers.RemoveAccidentalHits()
+
+    # This will zero the logits of negative candidates that have the same ID as
+    # the positive candidate from `labels` so as to not negatively impact the
+    # true positive.
+    logits = remove_accidental_hits(logits, labels, candidate_ids)
+    ```
     """
 
     def call(
@@ -29,16 +41,17 @@ class RemoveAccidentalHits(keras.layers.Layer):
         have the same ID as the positive candidate in that row.
 
         Args:
-            logits: logits tensor, typically `[batch_size, num_candidates]` but
-                can have more dimensions or be 1D as `[num_candidates]`.
-            labels: one-hot labels tensor, must be the same shape as `logits`.
-            candidate_ids: candidate identifiers tensor, can be
+            logits: The logits tensor, typically `[batch_size, num_candidates]`
+                but can have more dimensions or be 1D as `[num_candidates]`.
+            labels: The one-hot labels tensor, must be the same shape as
+                `logits`.
+            candidate_ids: The candidate identifiers tensor, can be
                 `[num_candidates]` or `[batch_size, num_candidates]` or have
                 more dimensions as long as they match the last dimensions of
                 `labels`.
 
         Returns:
-            logits: 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.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/keras_rs/src/layers/retrieval/sampling_probability_correction.py

@@ -17,6 +17,18 @@ class SamplingProbabilityCorrection(keras.layers.Layer):
         epsilon: float. Small float added to sampling probability to avoid
             taking the log of zero. Defaults to 1e-6.
         **kwargs: Args to pass to the base class.
+
+    Example:
+
+    ```python
+    # Create the layer.
+    sampling_probability_correction = (
+        keras_rs.layers.SamplingProbabilityCorrection()
+    )
+
+    # Correct the logits based on the provided candidate sampling probability.
+    logits = sampling_probability_correction(logits, probabilities)
+    ```
     """
 
     def __init__(self, epsilon: float = 1e-6, **kwargs: Any) -> None:
@@ -32,11 +44,14 @@ class SamplingProbabilityCorrection(keras.layers.Layer):
         """Corrects input logits to account for candidate sampling probability.
 
         Args:
-            logits: The logits to correct.
-            candidate_sampling_probability: The sampling probability.
+            logits: The logits tensor to correct, typically
+                `[batch_size, num_candidates]` but can have more dimensions or
+                be 1D as `[num_candidates]`.
+            candidate_sampling_probability: The sampling probability with the
+                same shape as `logits`.
 
         Returns:
-            The corrected logits.
+            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.dev2025043003/keras_rs/src/losses/pairwise_hinge_loss.py

@@ -0,0 +1,90 @@
+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.PairwiseHingeLoss")
+class PairwiseHingeLoss(PairwiseLoss):
+    def pairwise_loss(self, pairwise_logits: types.Tensor) -> types.Tensor:
+        return ops.relu(ops.subtract(ops.array(1), pairwise_logits))
+
+
+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`.
+"""
+extra_args = ""
+example = """
+    With `compile()` API:
+
+    ```python
+    model.compile(
+        loss=keras_rs.losses.PairwiseHingeLoss(),
+        ...
+    )
+    ```
+
+    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
+
+    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
+
+    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_hinge_loss(y_true=y_true, y_pred=y_pred)
+    0.64999
+
+    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_hinge_loss = keras_rs.losses.PairwiseHingeLoss()
+    >>> pairwise_hinge_loss(
+    ...     y_true=y_true, y_pred=y_pred, sample_weight=sample_weight
+    ... )
+    1.02499
+
+    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(
+    ...     reduction="none"
+    ... )
+    >>> pairwise_hinge_loss(y_true=y_true, y_pred=y_pred)
+    [[3. , 0. , 2. , 0.], [0., 0.20000005, 0.79999995, 0.]]
+"""
+
+PairwiseHingeLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="hinge loss",
+    formula=formula,
+    explanation=explanation,
+    extra_args=extra_args,
+    example=example,
+)

keras_rs_nightly-0.0.1.dev2025043003/keras_rs/src/losses/pairwise_logistic_loss.py

@@ -0,0 +1,99 @@
+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.PairwiseLogisticLoss")
+class PairwiseLogisticLoss(PairwiseLoss):
+    def pairwise_loss(self, pairwise_logits: types.Tensor) -> types.Tensor:
+        return ops.add(
+            ops.relu(ops.negative(pairwise_logits)),
+            ops.log(
+                ops.add(
+                    ops.array(1),
+                    ops.exp(ops.negative(ops.abs(pairwise_logits))),
+                )
+            ),
+        )
+
+
+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.
+"""
+extra_args = ""
+example = """
+    With `compile()` API:
+
+    ```python
+    model.compile(
+        loss=keras_rs.losses.PairwiseLogisticLoss(),
+        ...
+    )
+    ```
+
+    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
+
+    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
+
+    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_logistic_loss(y_true=y_true, y_pred=y_pred)
+    0.53751
+
+    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_logistic_loss = keras_rs.losses.PairwiseLogisticLoss()
+    >>> pairwise_logistic_loss(
+    ...     y_true=y_true, y_pred=y_pred, sample_weight=sample_weight
+    ... )
+    >>> 0.80337
+
+    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(
+    ...     reduction="none"
+    ... )
+    >>> pairwise_logistic_loss(y_true=y_true, y_pred=y_pred)
+    [[2.126928, 0., 1.3132616, 0.48877698], [0., 0.20000005, 0.79999995, 0.]]
+"""
+
+PairwiseLogisticLoss.__doc__ = pairwise_loss_subclass_doc_string.format(
+    loss_name="logistic loss",
+    formula=formula,
+    explanation=explanation,
+    extra_args=extra_args,
+    example=example,
+)

{keras_rs_nightly-0.0.1.dev2025042803 → keras_rs_nightly-0.0.1.dev2025043003}/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
@@ -154,5 +159,7 @@ pairwise_loss_subclass_doc_string = (
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
-    """
+
+    Examples:
+{example}"""
 )