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

keras_rs/src/metrics/ranking_metrics_utils.py

@@ -0,0 +1,257 @@
+from typing import Callable
+
+import keras
+from keras import ops
+
+from keras_rs.src import types
+
+
+def get_shuffled_indices(
+    shape: types.Shape,
+    mask: types.Tensor | None = None,
+    shuffle_ties: bool = True,
+    seed: int | keras.random.SeedGenerator | None = None,
+) -> types.Tensor:
+    """Utility function for getting shuffled indices, with masked indices
+    pushed to the end.
+
+    Args:
+        shape: tuple. The shape of the tensor for which to generate
+            shuffled indices.
+        mask: An optional boolean tensor with the same shape as `shape`.
+            If provided, elements where `mask` is `False` will be placed
+            at the end of the sorted indices. Defaults to `None` (no masking).
+        shuffle_ties: Boolean indicating how to handle ties if multiple elements
+            have the same sorting value (randomly when `shuffle_ties` is True
+            otherwise, order is preserved).
+        seed: Optional integer seed for the random number generator used when
+            `shuffle_ties` is True. Ensures reproducibility. Defaults to None.
+
+    Returns:
+        A tensor of shape `shape` containing shuffled indices.
+    """
+    # If `shuffle_ties` is True, generate random values. Otherwise, generate
+    # zeros so that we get `[0, 1, 2, ...]` as indices on doing `argsort`.
+    if shuffle_ties:
+        shuffle_values = keras.random.uniform(shape, seed=seed, dtype="float32")
+    else:
+        shuffle_values = ops.zeros(shape, dtype="float32")
+
+    # When `mask = False`, increase value by 1 so that those indices are placed
+    # at the end. Note that `shuffle_values` lies in the range `[0, 1)`, so
+    # adding by 1 works out.
+    if mask is not None:
+        shuffle_values = ops.where(
+            mask,
+            shuffle_values,
+            ops.add(shuffle_values, ops.cast(1, dtype="float32")),
+        )
+
+    shuffled_indices = ops.argsort(shuffle_values)
+    return shuffled_indices
+
+
+def sort_by_scores(
+    tensors_to_sort: list[types.Tensor],
+    scores: types.Tensor,
+    mask: types.Tensor | None = None,
+    k: int | None = None,
+    shuffle_ties: bool = True,
+    seed: int | keras.random.SeedGenerator | None = None,
+) -> types.Tensor:
+    """
+    Utility function for sorting tensors by scores.
+
+    Args:
+        tensors_to_sort: list of tensors. All tensors are of shape
+            `(batch_size, list_size)`. These tensors are sorted based on
+            `scores`.
+        scores: tensor. Of shape `(batch_size, list_size)`. The scores to sort
+            by.
+        k: int. The number of top-ranked items to consider (the 'k' in 'top-k').
+            If `None`, `list_size` is used.
+        shuffle_ties: bool. Whether to randomly shuffle scores before sorting.
+            This is done to break ties.
+        seed: int. Seed for shuffling.
+
+    Returns:
+        List of sorted tensors (`tensors_to_sort`), sorted using `scores`.
+    """
+    max_possible_k = ops.shape(scores)[1]
+    if k is None:
+        k = max_possible_k
+    elif isinstance(max_possible_k, int):
+        k = min(k, max_possible_k)
+    else:
+        k = ops.minimum(k, max_possible_k)
+
+    # --- Work around for PyTorch instability ---
+    # Torch's `topk` is not stable with `sorted=True`, unlike JAX and TF.
+    # See:
+    #   - https://github.com/pytorch/pytorch/issues/27542
+    #   - https://github.com/pytorch/pytorch/issues/88227
+    #
+    # This small "stable offset" ensures deterministic tie-breaking for
+    # equal scores. We can remove this workaround once PyTorch adds a
+    # `stable=True` flag for topk.
+
+    if keras.backend.backend() == "torch" and not shuffle_ties:
+        list_size = ops.shape(scores)[1]
+        indices = ops.arange(list_size)
+        indices = ops.expand_dims(indices, axis=0)
+        indices = ops.broadcast_to(indices, ops.shape(scores))
+        stable_offset = ops.cast(indices, scores.dtype) * 1e-6
+        scores = ops.subtract(scores, stable_offset)
+    # --- End FIX ---
+
+    # Shuffle ties randomly, and push masked values to the beginning.
+    shuffled_indices = None
+    if shuffle_ties or mask is not None:
+        shuffled_indices = get_shuffled_indices(
+            ops.shape(scores),
+            mask=mask,
+            shuffle_ties=True,
+            seed=seed,
+        )
+        scores = ops.take_along_axis(scores, shuffled_indices, axis=1)
+
+    # Get top-k indices.
+    _, indices = ops.top_k(scores, k=k, sorted=True)
+
+    # If we shuffled our `scores` tensor, we need to get the correct indices
+    # by indexing into `shuffled_indices`.
+    if shuffled_indices is not None:
+        indices = ops.take_along_axis(shuffled_indices, indices, axis=1)
+
+    return [
+        ops.take_along_axis(tensor_to_sort, indices, axis=1)
+        for tensor_to_sort in tensors_to_sort
+    ]
+
+
+def get_list_weights(
+    weights: types.Tensor, relevance: types.Tensor
+) -> types.Tensor:
+    """Computes per-list weights from provided sample weights.
+
+    Per-list weights are calculated as follows:
+    ```
+    per_list_weights = sum(weights * relevance) / sum(relevance).
+    ```
+
+    For lists where the sum of relevance is 0, a default weight is assigned:
+    ```
+    sum(per_list_weights) / num(sum(relevance) != 0 AND sum(weights) != 0)
+    ```
+
+    If all lists have a sum of relevance equal to 0, the default weight is 1.0.
+
+    As a result of the above computation, this function takes care of the
+    following cases:
+    - **Uniform Weights:** When all input weights are 1.0, all per-list weights
+      will be 1.0, even for lists with no relevant examples. This aligns with
+      standard ranking metrics.
+    - **Non-zero Weights per List:** If every list has at least one non-zero
+      weight, the default weight mechanism is not utilized, which is suitable
+      for unbiased metrics.
+    - **Mixed Scenarios:** For cases with a mix of lists having zero and
+      non-zero relevance and weights, the weights for lists with non-zero
+      relevance and weights are proportional to:
+
+      ```
+      per_list_weights / sum(per_list_weights) *
+      num(sum(relevance) != 0) / num(lists)
+      ```
+
+      The rest have weights `1.0 / num(lists)`.
+
+    Args:
+        weights: tensor. Weights tensor of shape `(batch_size, list_size)`.
+        relevance: tensor. The relevance `Tensor` of shape
+            `(batch_size, list_size)`.
+
+    Returns:
+        A tensor of shape `(batch_size, 1)`, containing the per-list weights.
+    """
+    # Calculate if the sum of weights per list is greater than 0.0.
+    nonzero_weights = ops.greater(ops.sum(weights, axis=1, keepdims=True), 0.0)
+    # Calculate the sum of relevance per list
+    per_list_relevance = ops.sum(relevance, axis=1, keepdims=True)
+    # Calculate if the sum of relevance per list is greater than 0.0
+    nonzero_relevance_condition = ops.greater(per_list_relevance, 0.0)
+    # Identify lists where both weights and relevance sums are non-zero.
+    nonzero_relevance = ops.cast(
+        ops.logical_and(nonzero_weights, nonzero_relevance_condition),
+        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)
+
+    # Calculate the per-list weights using the core formula.
+    # Numerator: `sum(weights * relevance)` per list
+    numerator = ops.sum(ops.multiply(weights, relevance), axis=1, keepdims=True)
+    # Denominator: per_list_relevance = sum(relevance) per list
+    per_list_weights = ops.divide_no_nan(numerator, per_list_relevance)
+
+    # Calculate the sum of the computed per-list weights.
+    sum_weights = ops.sum(per_list_weights, axis=0, keepdims=True)
+
+    # Calculate the average weight to use as default for lists with zero
+    # relevance but non-zero weights. If no lists have non-zero relevance,
+    # default to 1.0.
+    avg_weight = ops.where(
+        ops.greater(nonzero_relevance_count, 0.0),
+        ops.divide(sum_weights, nonzero_relevance_count),
+        ops.cast(1, dtype=sum_weights.dtype),
+    )
+
+    # Final assignment of weights based on conditions:
+    # 1. If sum(weights) == 0 for a list, the final weight is 0.
+    # 2. If sum(weights) > 0 AND sum(relevance) > 0, use the calculated
+    # `per_list_weights`.
+    # 3. If `sum(weights) > 0` AND `sum(relevance) == 0`, use the calculated
+    # `avg_weight`.
+    final_weights = ops.where(
+        nonzero_weights,
+        ops.where(
+            nonzero_relevance_condition,
+            per_list_weights,
+            avg_weight,
+        ),
+        ops.cast(0, dtype=per_list_weights.dtype),
+    )
+
+    return final_weights
+
+
+@keras.saving.register_keras_serializable()  # type: ignore[untyped-decorator]
+def default_gain_fn(label: types.Tensor) -> types.Tensor:
+    return ops.subtract(ops.power(2.0, label), 1.0)
+
+
+@keras.saving.register_keras_serializable()  # type: ignore[untyped-decorator]
+def default_rank_discount_fn(rank: types.Tensor) -> types.Tensor:
+    return ops.divide(
+        ops.cast(1, dtype=rank.dtype),
+        ops.log2(ops.add(ops.cast(1, dtype=rank.dtype), rank)),
+    )
+
+
+def compute_dcg(
+    y_true: types.Tensor,
+    sample_weight: types.Tensor,
+    gain_fn: Callable[[types.Tensor], types.Tensor] = default_gain_fn,
+    rank_discount_fn: Callable[
+        [types.Tensor], types.Tensor
+    ] = default_rank_discount_fn,
+) -> types.Tensor:
+    list_size = ops.shape(y_true)[1]
+    positions = ops.arange(1, list_size + 1, dtype=y_true.dtype)
+    gain = gain_fn(y_true)
+    discount = rank_discount_fn(positions)
+
+    return ops.sum(
+        ops.multiply(sample_weight, ops.multiply(gain, discount)),
+        axis=1,
+        keepdims=True,
+    )

keras_rs/src/metrics/recall_at_k.py

@@ -0,0 +1,108 @@
+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.RecallAtK")
+class RecallAtK(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,
+        )
+
+        relevance = ops.cast(
+            ops.greater_equal(
+                sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
+            ),
+            dtype=y_pred.dtype,
+        )
+        overall_relevance = ops.cast(
+            ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
+            dtype=y_pred.dtype,
+        )
+        per_list_recall = ops.divide_no_nan(
+            ops.sum(relevance, axis=1, keepdims=True),
+            ops.sum(overall_relevance, axis=1, keepdims=True),
+        )
+
+        # Get weights.
+        per_list_weights = get_list_weights(
+            weights=sample_weight, relevance=overall_relevance
+        )
+
+        return per_list_recall, per_list_weights
+
+
+concept_sentence = (
+    "It measures the proportion of relevant items found in the top-k "
+    "recommendations out of the total number of relevant items for a user"
+)
+relevance_type = "binary indicators (0 or 1) of relevance"
+score_range_interpretation = (
+    "Scores range from 0 to 1, with 1 indicating that all relevant items "
+    "for the user were found within the top-k recommendations"
+)
+formula = """```
+R@k(y, s) = sum_i I[rank(s_i) < k] y_i / sum_j y_j
+```
+
+where `y_i` is the relevance label (0/1) of the item ranked at position
+`i`, `I[condition]` is 1 if the condition is met, otherwise 0."""
+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.RecallAtK()(
+    ...     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.RecallAtK()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+RecallAtK.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Recall@k",
+    metric_abbreviation="R@k",
+    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/utils.py

@@ -0,0 +1,70 @@
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.utils.keras_utils import check_rank
+from keras_rs.src.utils.keras_utils import check_shapes_compatible
+
+
+def standardize_call_inputs_ranks(
+    y_true: types.Tensor,
+    y_pred: types.Tensor,
+    mask: types.Tensor | None = None,
+    check_y_true_rank: bool = True,
+) -> tuple[types.Tensor, types.Tensor, types.Tensor | None, bool]:
+    """
+    Utility function for processing inputs for losses and metrics.
+
+    This utility function does three things:
+
+    - Checks that `y_true`, `y_pred` are of rank 1 or 2;
+    - Checks that `y_true`, `y_pred`, `mask` have the same shape;
+    - Adds batch dimension if rank = 1.
+
+    Args:
+        y_true: tensor. Ground truth values.
+        y_pred: tensor. The predicted values.
+        mask: tensor. Boolean mask for `y_true`.
+        check_y_true_rank: bool. Whether to check the rank of `y_true`.
+
+    Returns:
+        Tuple of processed `y_true`, `y_pred`, `mask`, and `batched`. `batched`
+        is a bool indicating if the inputs are batched.
+    """
+
+    y_true_shape = ops.shape(y_true)
+    y_true_rank = len(y_true_shape)
+    y_pred_shape = ops.shape(y_pred)
+    y_pred_rank = len(y_pred_shape)
+    if mask is not None:
+        mask_shape = ops.shape(mask)
+        mask_rank = len(mask_shape)
+
+    if check_y_true_rank:
+        check_rank(y_true_rank, allowed_ranks=(1, 2), tensor_name="y_true")
+    check_rank(y_pred_rank, allowed_ranks=(1, 2), tensor_name="y_pred")
+    if mask is not None:
+        check_rank(mask_rank, allowed_ranks=(1, 2), tensor_name="mask")
+    if not check_shapes_compatible(y_true_shape, y_pred_shape):
+        raise ValueError(
+            "`y_true` and `y_pred` should have the same shape. Received: "
+            f"`y_true.shape` = {y_true_shape}, `y_pred.shape` = {y_pred_shape}."
+        )
+    if mask is not None and not check_shapes_compatible(
+        y_true_shape, mask_shape
+    ):
+        raise ValueError(
+            "`y_true['labels']` and `y_true['mask']` should have the same "
+            f"shape. Received: `y_true['labels'].shape` = {y_true_shape}, "
+            f"`y_true['mask'].shape` = {mask_shape}."
+        )
+
+    batched = True
+    if y_true_rank == 1:
+        batched = False
+
+        y_true = ops.expand_dims(y_true, axis=0)
+        y_pred = ops.expand_dims(y_pred, axis=0)
+        if mask is not None:
+            mask = ops.expand_dims(mask, axis=0)
+
+    return y_true, y_pred, mask, batched

keras_rs/src/types.py

@@ -1,6 +1,8 @@
 """Type definitions."""
 
-from typing import Any, Optional, Sequence
+from typing import Any, Callable, Mapping, Sequence, TypeAlias, TypeVar, Union
+
+import keras
 
 """
 A tensor in any of the backends.
@@ -8,19 +10,46 @@ A tensor in any of the backends.
 We do not define it explicitly to not require all the backends to be installed
 and imported. The explicit definition would be:
 ```
-Union[
-  numpy.ndarray,
-  tensorflow.Tensor,
-  tensorflow.RaggedTensor,
-  tensorflow.SparseTensor,
-  tensorflow.IndexedSlices,
-  jax.Array,
-  jax.experimental.sparse.JAXSparse,
-  torch.Tensor,
-  keras.KerasTensor,
-]
+numpy.ndarray,
+| tensorflow.Tensor,
+| tensorflow.RaggedTensor,
+| tensorflow.SparseTensor,
+| tensorflow.IndexedSlices,
+| jax.Array,
+| jax.experimental.sparse.JAXSparse,
+| torch.Tensor,
+| keras.KerasTensor,
 ```
 """
-Tensor = Any
+Tensor: TypeAlias = Any
+
+Shape: TypeAlias = Sequence[int | None]
+
+DType: TypeAlias = str
+
+ConstraintLike: TypeAlias = (
+    str
+    | keras.constraints.Constraint
+    | type[keras.constraints.Constraint]
+    | Callable[[Tensor], Tensor]
+)
+
+InitializerLike: TypeAlias = (
+    str
+    | keras.initializers.Initializer
+    | type[keras.initializers.Initializer]
+    | Callable[[Shape, DType], Tensor]
+    | Tensor
+)
+
+RegularizerLike: TypeAlias = (
+    str
+    | keras.regularizers.Regularizer
+    | type[keras.regularizers.Regularizer]
+    | Callable[[Tensor], Tensor]
+)
 
-TensorShape = Sequence[Optional[int]]
+T = TypeVar("T")
+Nested: TypeAlias = (
+    T | Sequence[Union[T, "Nested[T]"]] | Mapping[str, Union[T, "Nested[T]"]]
+)

keras_rs/src/utils/doc_string_utils.py

@@ -0,0 +1,53 @@
+import re
+import textwrap
+from typing import Any
+
+
+def format_docstring(template: str, width: int = 80, **kwargs: Any) -> str:
+    """Formats and wraps a docstring using dedent and fill."""
+    base_indent_str = " " * 4
+
+    # Initial format
+    formatted = template.format(**kwargs)
+
+    # Dedent the whole block
+    dedented_all = textwrap.dedent(formatted).strip()
+
+    # Split into logical paragraphs/blocks.
+    blocks = re.split(r"(\n\s*\n)", dedented_all)
+
+    processed_output = []
+
+    for block in blocks:
+        stripped_block = block.strip()
+        if not stripped_block:
+            processed_output.append(block)
+            continue
+
+        if "```" in stripped_block:
+            formula_dedented = textwrap.dedent(stripped_block)
+            processed_output.append(
+                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(
+                    splitted_block[0] + "\n\n" + "\n".join(splitted_block[1:]),
+                    base_indent_str,
+                )
+            )
+        else:
+            processed_output.append(
+                textwrap.fill(
+                    stripped_block,
+                    width=width - len(base_indent_str),
+                    initial_indent=base_indent_str,
+                    subsequent_indent=base_indent_str,
+                )
+            )
+
+    final_string = "".join(processed_output).strip()
+    final_string = base_indent_str + final_string
+    return final_string

keras_rs/src/utils/keras_utils.py

@@ -1,11 +1,35 @@
-from typing import Union
+from typing import Any, Callable
 
 import keras
 
+from keras_rs.src import types
+
+
+def no_automatic_dependency_tracking(
+    fn: Callable[..., Any],
+) -> Callable[..., Any]:
+    """Decorator to disable automatic dependency tracking in Keras and TF.
+
+    Args:
+        fn: the function to disable automatic dependency tracking for.
+
+    Returns:
+        a wrapped version of `fn`.
+    """
+    if keras.backend.backend() == "tensorflow":
+        import tensorflow as tf
+
+        fn = tf.__internal__.tracking.no_automatic_dependency_tracking(fn)
+
+    wrapped_fn: Callable[..., Any] = (
+        keras.src.utils.tracking.no_automatic_dependency_tracking(fn)
+    )
+    return wrapped_fn
+
 
 def clone_initializer(
-    initializer: Union[str, keras.initializers.Initializer],
-) -> keras.initializers.Initializer:
+    initializer: types.InitializerLike,
+) -> types.InitializerLike:
     """Clones an initializer to ensure a new seed.
 
     Args:
@@ -25,3 +49,28 @@ def clone_initializer(
         return initializer_class.from_config(config)
     # If we get a string or dict, just return as we cannot and should not clone.
     return initializer
+
+
+def check_shapes_compatible(shape1: types.Shape, shape2: types.Shape) -> bool:
+    # Check rank first.
+    if len(shape1) != len(shape2):
+        return False
+
+    for d1, d2 in zip(shape1, shape2):
+        if isinstance(d1, int) and isinstance(d2, int):
+            if d1 != d2:
+                return False
+
+    return True
+
+
+def check_rank(
+    x_rank: int,
+    allowed_ranks: tuple[int, ...],
+    tensor_name: str,
+) -> None:
+    if x_rank not in allowed_ranks:
+        raise ValueError(
+            f"`{tensor_name}` should have a rank from `{allowed_ranks}`."
+            f"Received: `{x_rank}`."
+        )