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

keras_rs/src/layers/{modeling → feature_interaction}/feature_cross.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Text, Union
+from typing import Any
 
 import keras
 from keras import ops
@@ -52,17 +52,37 @@ class FeatureCross(keras.layers.Layer):
             Regularizer to use for the kernel matrix.
         bias_regularizer: string or `keras.regularizer` regularizer.
             Regularizer to use for the bias vector.
+        **kwargs: Args to pass to the base class.
 
     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:
@@ -72,20 +92,18 @@ class FeatureCross(keras.layers.Layer):
 
     def __init__(
         self,
-        projection_dim: Optional[int] = None,
-        diag_scale: Optional[float] = 0.0,
+        projection_dim: int | None = None,
+        diag_scale: float | None = 0.0,
         use_bias: bool = True,
-        pre_activation: Optional[Union[str, keras.layers.Activation]] = None,
-        kernel_initializer: Union[
-            Text, keras.initializers.Initializer
-        ] = "glorot_uniform",
-        bias_initializer: Union[Text, keras.initializers.Initializer] = "zeros",
-        kernel_regularizer: Union[
-            Text, None, keras.regularizers.Regularizer
-        ] = None,
-        bias_regularizer: Union[
-            Text, None, keras.regularizers.Regularizer
-        ] = None,
+        pre_activation: str | keras.layers.Activation | None = None,
+        kernel_initializer: (
+            str | keras.initializers.Initializer
+        ) = "glorot_uniform",
+        bias_initializer: str | keras.initializers.Initializer = "zeros",
+        kernel_regularizer: (
+            str | None | keras.regularizers.Regularizer
+        ) = None,
+        bias_regularizer: (str | None | keras.regularizers.Regularizer) = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -109,7 +127,7 @@ class FeatureCross(keras.layers.Layer):
                 f"`diag_scale={self.diag_scale}`"
             )
 
-    def build(self, input_shape: types.TensorShape) -> None:
+    def build(self, input_shape: types.Shape) -> None:
         last_dim = input_shape[-1]
 
         if self.projection_dim is not None:
@@ -135,7 +153,7 @@ class FeatureCross(keras.layers.Layer):
         self.built = True
 
     def call(
-        self, x0: types.Tensor, x: Optional[types.Tensor] = None
+        self, x0: types.Tensor, x: types.Tensor | None = None
     ) -> types.Tensor:
         """Forward pass of the cross layer.
 

keras_rs/src/layers/retrieval/brute_force_retrieval.py

@@ -1,13 +1,14 @@
-from typing import Any, Optional, Union
+from typing import Any
 
 import keras
 
 from keras_rs.src import types
 from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.layers.retrieval.retrieval import Retrieval
 
 
 @keras_rs_export("keras_rs.layers.BruteForceRetrieval")
-class BruteForceRetrieval(keras.layers.Layer):
+class BruteForceRetrieval(Retrieval):
     """Brute force top-k retrieval.
 
     This layer maintains a set of candidates and is able to exactly retrieve the
@@ -54,17 +55,19 @@ class BruteForceRetrieval(keras.layers.Layer):
 
     def __init__(
         self,
-        candidate_embeddings: Optional[types.Tensor] = None,
-        candidate_ids: Optional[types.Tensor] = None,
+        candidate_embeddings: types.Tensor | None = None,
+        candidate_ids: types.Tensor | None = None,
         k: int = 10,
         return_scores: bool = True,
         **kwargs: Any,
     ) -> None:
-        super().__init__(**kwargs)
+        # Keep `k`, `return_scores` as separately passed args instead of keeping
+        # them in `kwargs`. This is to ensure the user does not have to hop
+        # to the base class to check which other args can be passed.
+        super().__init__(k=k, return_scores=return_scores, **kwargs)
+
         self.candidate_embeddings = None
         self.candidate_ids = None
-        self.k = k
-        self.return_scores = return_scores
 
         if candidate_embeddings is None:
             if candidate_ids is not None:
@@ -78,42 +81,18 @@ class BruteForceRetrieval(keras.layers.Layer):
     def update_candidates(
         self,
         candidate_embeddings: types.Tensor,
-        candidate_ids: Optional[types.Tensor] = None,
+        candidate_ids: types.Tensor | None = None,
     ) -> None:
         """Update the set of candidates and optionally their candidate IDs.
 
         Args:
             candidate_embeddings: The candidate embeddings.
-            candidate_ids: The identifiers for the candidates. If `None` the
+            candidate_ids: The identifiers for the candidates. If `None`, the
                 indices of the candidates are returned instead.
         """
-        if candidate_embeddings is None:
-            raise ValueError("`candidate_embeddings` is required")
-
-        if len(candidate_embeddings.shape) != 2:
-            raise ValueError(
-                "`candidate_embeddings` must be a tensor of rank 2 "
-                "(num_candidates, embedding_size), received "
-                "`candidate_embeddings` with shape "
-                f"{candidate_embeddings.shape}"
-            )
-
-        if candidate_embeddings.shape[0] < self.k:
-            raise ValueError(
-                "The number of candidates provided "
-                f"({candidate_embeddings.shape[0]}) is less than the number of "
-                f"candidates to retrieve (k={self.k})."
-            )
-
-        if (
-            candidate_ids is not None
-            and candidate_ids.shape[0] != candidate_embeddings.shape[0]
-        ):
-            raise ValueError(
-                "The `candidate_embeddings` and `candidate_is` tensors must "
-                "have the same number of rows, got tensors of shape "
-                f"{candidate_embeddings.shape} and {candidate_ids.shape}."
-            )
+        self._validate_candidate_embeddings_and_ids(
+            candidate_embeddings, candidate_ids
+        )
 
         if self.candidate_embeddings is not None:
             # Update of existing variables.
@@ -146,7 +125,7 @@ class BruteForceRetrieval(keras.layers.Layer):
 
     def call(
         self, inputs: types.Tensor
-    ) -> Union[types.Tensor, tuple[types.Tensor, types.Tensor]]:
+    ) -> types.Tensor | tuple[types.Tensor, types.Tensor]:
         """Returns the top candidates for the query passed as input.
 
         Args:
@@ -167,31 +146,3 @@ class BruteForceRetrieval(keras.layers.Layer):
             return top_scores, top_ids
         else:
             return top_ids
-
-    def compute_score(
-        self, query_embedding: types.Tensor, candidate_embedding: types.Tensor
-    ) -> types.Tensor:
-        """Computes the standard dot product score from queries and candidates.
-
-        Args:
-            query_embedding: Tensor of query embedding corresponding to the
-                queries for which to retrieve top candidates.
-            candidate_embedding: Tensor of candidate embeddings.
-
-        Returns:
-            The dot product of queries and candidates.
-        """
-
-        return keras.ops.matmul(
-            query_embedding, keras.ops.transpose(candidate_embedding)
-        )
-
-    def get_config(self) -> dict[str, Any]:
-        config: dict[str, Any] = super().get_config()
-        config.update(
-            {
-                "k": self.k,
-                "return_scores": self.compute_score,
-            }
-        )
-        return config

keras_rs/src/layers/retrieval/hard_negative_mining.py

@@ -0,0 +1,94 @@
+from typing import Any
+
+import keras
+import ml_dtypes
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+
+MAX_FLOAT = ml_dtypes.finfo("float32").max / 100.0
+
+
+@keras_rs_export("keras_rs.layers.HardNegativeMining")
+class HardNegativeMining(keras.layers.Layer):
+    """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:
+        super().__init__(**kwargs)
+        self._num_hard_negatives = num_hard_negatives
+        self.built = True
+
+    def call(
+        self, logits: types.Tensor, labels: types.Tensor
+    ) -> tuple[types.Tensor, types.Tensor]:
+        """Filters logits and labels with per-query hard negative mining.
+
+        The result will include logits and labels for `num_hard_negatives`
+        negatives as well as the positive candidate.
+
+        Args:
+            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:
+            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.
+        """
+
+        # Number of sampled logits, i.e, the number of hard negatives to be
+        # sampled (k) + number of true logit (1) per query, capped by batch
+        # size.
+        num_logits = ops.shape(logits)[-1]
+        if isinstance(num_logits, int):
+            num_sampled = min(self._num_hard_negatives + 1, num_logits)
+        else:
+            num_sampled = ops.minimum(self._num_hard_negatives + 1, num_logits)
+        # To gather indices of top k negative logits per row (query) in logits,
+        # true logits need to be excluded. First replace the true logits
+        # (corresponding to positive labels) with a large score value and then
+        # select the top k + 1 logits from each row so that selected indices
+        # include the indices of true logit + top k negative logits. This
+        # approach is to avoid using inefficient masking when excluding true
+        # logits.
+
+        # For each query, get the indices of the logits which have the highest
+        # k + 1 logit values, including the highest k negative logits and one
+        # true logit.
+        _, indices = ops.top_k(
+            ops.add(logits, ops.multiply(labels, MAX_FLOAT)),
+            k=num_sampled,
+            sorted=False,
+        )
+
+        # Gather sampled logits and corresponding labels.
+        logits = ops.take_along_axis(logits, indices, axis=-1)
+        labels = ops.take_along_axis(labels, indices, axis=-1)
+
+        return logits, labels

keras_rs/src/layers/retrieval/remove_accidental_hits.py

@@ -0,0 +1,97 @@
+import keras
+import ml_dtypes
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.utils import keras_utils
+
+SMALLEST_FLOAT = ml_dtypes.finfo("float32").smallest_normal / 100.0
+
+
+@keras_rs_export("keras_rs.layers.RemoveAccidentalHits")
+class RemoveAccidentalHits(keras.layers.Layer):
+    """Zeroes the logits of accidental negatives.
+
+    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(
+        self,
+        logits: types.Tensor,
+        labels: types.Tensor,
+        candidate_ids: types.Tensor,
+    ) -> types.Tensor:
+        """Zeroes selected logits.
+
+        For each row in the batch, zeroes the logits of negative candidates that
+        have the same ID as the positive candidate in that row.
+
+        Args:
+            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:
+            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
+        # approximate so by letting accidental hits have extremely small logits
+        # (SMALLEST_FLOAT) for ease-of-implementation.
+
+        labels_shape = ops.shape(labels)
+        labels_rank = len(labels_shape)
+        logits_shape = ops.shape(logits)
+        candidate_ids_shape = ops.shape(candidate_ids)
+        candidate_ids_rank = len(candidate_ids_shape)
+
+        if not keras_utils.check_shapes_compatible(labels_shape, logits_shape):
+            raise ValueError(
+                "`labels` and `logits` should have the same shape. Received: "
+                f"`labels.shape` = {labels_shape}, "
+                f"`logits.shape` = {logits_shape}."
+            )
+
+        if not keras_utils.check_shapes_compatible(
+            labels_shape[-candidate_ids_rank:], candidate_ids_shape
+        ):
+            raise ValueError(
+                "`candidate_ids` should have the same shape as the last "
+                "dimensions of `labels`. Received: "
+                f"`candidate_ids.shape` = {candidate_ids_shape}, "
+                f"`labels.shape` = {labels_shape}."
+            )
+
+        # Add dimensions to `candidate_ids` to have the same rank as `labels`.
+        if candidate_ids_rank < labels_rank:
+            candidate_ids = ops.expand_dims(
+                candidate_ids, list(range(labels_rank - candidate_ids_rank))
+            )
+        positive_indices = ops.expand_dims(ops.argmax(labels, axis=-1), -1)
+        positive_candidate_ids = ops.take(candidate_ids, positive_indices)
+
+        duplicate = ops.cast(
+            ops.equal(positive_candidate_ids, candidate_ids), labels.dtype
+        )
+        duplicate = ops.subtract(duplicate, labels)
+
+        return ops.add(logits, ops.multiply(duplicate, SMALLEST_FLOAT))

keras_rs/src/layers/retrieval/retrieval.py

@@ -0,0 +1,127 @@
+import abc
+from typing import Any
+
+import keras
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+
+
+@keras_rs_export("keras_rs.layers.Retrieval")
+class Retrieval(keras.layers.Layer, abc.ABC):
+    """Retrieval base abstract class.
+
+    This layer provides a common interface for all retrieval layers. In order
+    to implement a custom retrieval layer, this abstract class should be
+    subclassed.
+
+    Args:
+        k: int. Number of candidates to retrieve.
+        return_scores: bool. When `True`, this layer returns a tuple with the
+            top scores and the top identifiers. When `False`, this layer returns
+            a single tensor with the top identifiers.
+    """
+
+    def __init__(
+        self,
+        k: int = 10,
+        return_scores: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.k = k
+        self.return_scores = return_scores
+
+    def _validate_candidate_embeddings_and_ids(
+        self,
+        candidate_embeddings: types.Tensor,
+        candidate_ids: types.Tensor | None = None,
+    ) -> None:
+        """Validates inputs to `update_candidates()`."""
+
+        if candidate_embeddings is None:
+            raise ValueError("`candidate_embeddings` is required.")
+
+        if len(candidate_embeddings.shape) != 2:
+            raise ValueError(
+                "`candidate_embeddings` must be a tensor of rank 2 "
+                "(num_candidates, embedding_size), received "
+                "`candidate_embeddings` with shape "
+                f"{candidate_embeddings.shape}"
+            )
+
+        if candidate_embeddings.shape[0] < self.k:
+            raise ValueError(
+                "The number of candidates provided "
+                f"({candidate_embeddings.shape[0]}) is less than the number of "
+                f"candidates to retrieve (k={self.k})."
+            )
+
+        if (
+            candidate_ids is not None
+            and candidate_ids.shape[0] != candidate_embeddings.shape[0]
+        ):
+            raise ValueError(
+                "The `candidate_embeddings` and `candidate_is` tensors must "
+                "have the same number of rows, got tensors of shape "
+                f"{candidate_embeddings.shape} and {candidate_ids.shape}."
+            )
+
+    @abc.abstractmethod
+    def update_candidates(
+        self,
+        candidate_embeddings: types.Tensor,
+        candidate_ids: types.Tensor | None = None,
+    ) -> None:
+        """Update the set of candidates and optionally their candidate IDs.
+
+        Args:
+            candidate_embeddings: The candidate embeddings.
+            candidate_ids: The identifiers for the candidates. If `None`, the
+                indices of the candidates are returned instead.
+        """
+        pass
+
+    @abc.abstractmethod
+    def call(
+        self, inputs: types.Tensor
+    ) -> types.Tensor | tuple[types.Tensor, types.Tensor]:
+        """Returns the top candidates for the query passed as input.
+
+        Args:
+            inputs: the query for which to return top candidates.
+
+        Returns:
+            A tuple with the top scores and the top identifiers if
+            `returns_scores` is True, otherwise a tensor with the top
+            identifiers.
+        """
+        pass
+
+    def compute_score(
+        self, query_embedding: types.Tensor, candidate_embedding: types.Tensor
+    ) -> types.Tensor:
+        """Computes the standard dot product score from queries and candidates.
+
+        Args:
+            query_embedding: Tensor of query embedding corresponding to the
+                queries for which to retrieve top candidates.
+            candidate_embedding: Tensor of candidate embeddings.
+
+        Returns:
+            The dot product of queries and candidates.
+        """
+
+        return keras.ops.matmul(
+            query_embedding, keras.ops.transpose(candidate_embedding)
+        )
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update(
+            {
+                "k": self.k,
+                "return_scores": self.compute_score,
+            }
+        )
+        return config

keras_rs/src/layers/retrieval/sampling_probability_correction.py

@@ -0,0 +1,63 @@
+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
+
+
+@keras_rs_export("keras_rs.layers.SamplingProbabilityCorrection")
+class SamplingProbabilityCorrection(keras.layers.Layer):
+    """Sampling probability correction.
+
+    Corrects the logits to reflect the sampling probability of negatives.
+
+    Args:
+        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:
+        super().__init__(**kwargs)
+        self.epsilon = epsilon
+        self.built = True
+
+    def call(
+        self,
+        logits: types.Tensor,
+        candidate_sampling_probability: types.Tensor,
+    ) -> types.Tensor:
+        """Corrects input logits to account for candidate sampling probability.
+
+        Args:
+            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 with the same shape as the input logits.
+        """
+        return logits - ops.log(
+            ops.clip(candidate_sampling_probability, self.epsilon, 1.0)
+        )
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update({"epsilon": self.epsilon})
+        return config