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

keras_rs/src/losses/list_mle_loss.py

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

keras_rs/src/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/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/src/losses/pairwise_loss.py

@@ -0,0 +1,165 @@
+import abc
+from typing import Any
+
+import keras
+from keras import ops
+
+from keras_rs.src import types
+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):
+    """Base class for pairwise ranking losses.
+
+    Pairwise loss functions are designed for ranking tasks, where the goal is to
+    correctly order items within each list. Any pairwise loss function computes
+    the loss value by comparing pairs of items within each list, penalizing
+    cases where an item with a higher true label has a lower predicted score
+    than an item with a lower true label.
+
+    In order to implement any kind of pairwise loss, override the
+    `pairwise_loss` method.
+    """
+
+    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:
+        pass
+
+    def compute_unreduced_loss(
+        self,
+        labels: types.Tensor,
+        logits: types.Tensor,
+        mask: types.Tensor | None = None,
+    ) -> tuple[types.Tensor, types.Tensor]:
+        # 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)
+
+        pairwise_labels, pairwise_logits = pairwise_comparison(
+            labels=labels,
+            logits=logits,
+            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
+
+    def call(
+        self,
+        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)`
+                for batched inputs. If an item has a label of -1, it is ignored
+                in loss computation. If it is a dictionary, it should have two
+                keys: `"labels"` and `"mask"`. `"mask"` can be used to ignore
+                elements in loss computation, i.e., pairs will not be formed
+                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`.
+
+        Returns:
+            The loss.
+        """
+        mask = None
+        if isinstance(y_true, dict):
+            if "labels" not in y_true:
+                raise ValueError(
+                    '`"labels"` should be present in `y_true`. Received: '
+                    f"`y_true` = {y_true}"
+                )
+
+            mask = y_true.get("mask", None)
+            y_true = y_true["labels"]
+
+        y_true = ops.convert_to_tensor(y_true)
+        y_pred = ops.convert_to_tensor(y_pred)
+        if mask is not None:
+            mask = ops.convert_to_tensor(mask)
+
+        y_true, y_pred, mask, _ = standardize_call_inputs_ranks(
+            y_true, y_pred, mask
+        )
+
+        losses, weights = self.compute_unreduced_loss(
+            labels=y_true, logits=y_pred, mask=mask
+        )
+        losses = ops.multiply(losses, weights)
+        losses = ops.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 {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
+    pairs of items within each list, penalizing cases where an item with a
+    higher true label has a lower predicted score than an item with a lower
+    true label.
+
+    For each list of predicted scores `s` in `y_pred` and the corresponding list
+    of true labels `y` in `y_true`, the loss is computed as follows:
+
+    ```
+    {formula}
+    ```
+
+    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}
+    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
+            `"sum"`, `"sum_over_batch_size"`, `"mean"`,
+            `"mean_with_sample_weight"` or `None`. `"sum"` sums the loss,
+            `"sum_over_batch_size"` and `"mean"` sum the loss and divide by the
+            sample size, and `"mean_with_sample_weight"` sums the loss and
+            divides by the sum of the sample weights. `"none"` and `None`
+            perform no aggregation. Defaults to `"sum_over_batch_size"`.
+        name: Optional name for the loss instance.
+        dtype: The dtype of the loss's computations. Defaults to `None`, which
+            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
+            `"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}"""
+)

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