keras-rs-nightly 0.0.1.dev2025043003__py3-none-any.whl → 0.2.2.dev202506100336__py3-none-any.whl

keras_rs/src/layers/embedding/tensorflow/distributed_embedding.py

@@ -0,0 +1,424 @@
+from typing import Any, Callable, Sequence, TypeAlias
+
+import keras
+import tensorflow as tf
+
+from keras_rs.src import types
+from keras_rs.src.layers.embedding import base_distributed_embedding
+from keras_rs.src.layers.embedding import distributed_embedding_config
+from keras_rs.src.layers.embedding.tensorflow import config_conversion
+from keras_rs.src.utils import keras_utils
+
+FeatureConfig = distributed_embedding_config.FeatureConfig
+TableConfig = distributed_embedding_config.TableConfig
+
+# Placeholder of tf.tpu.experimental.embedding._Optimizer which is not exposed.
+TfTpuOptimizer: TypeAlias = Any
+
+
+GRADIENT_TRAP_DUMMY_NAME = "_gradient_trap_dummy"
+
+EMBEDDING_FEATURE_V1 = tf.tpu.experimental.HardwareFeature.EmbeddingFeature.V1
+EMBEDDING_FEATURE_V2 = tf.tpu.experimental.HardwareFeature.EmbeddingFeature.V2
+UNSUPPORTED = tf.tpu.experimental.HardwareFeature.EmbeddingFeature.UNSUPPORTED
+
+
+class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
+    """TensorFlow implementation of the TPU embedding layer."""
+
+    def __init__(
+        self,
+        feature_configs: types.Nested[
+            FeatureConfig | tf.tpu.experimental.embedding.FeatureConfig
+        ],
+        *,
+        table_stacking: (
+            str | Sequence[str] | Sequence[Sequence[str]]
+        ) = "auto",
+        **kwargs: Any,
+    ) -> None:
+        # Intercept arguments that are supported only on TensorFlow.
+        self._optimizer = kwargs.pop("optimizer", None)
+        self._pipeline_execution_with_tensor_core = kwargs.pop(
+            "pipeline_execution_with_tensor_core", False
+        )
+        self._sparse_core_embedding_config = kwargs.pop(
+            "sparse_core_embedding_config", None
+        )
+
+        # Mark as True by default for `_verify_input_shapes`. This will be
+        # updated in `_sparsecore_init` if applicable.
+        self._using_keras_rs_configuration = True
+
+        super().__init__(
+            feature_configs, table_stacking=table_stacking, **kwargs
+        )
+
+    def _is_tpu_strategy(self, strategy: tf.distribute.Strategy) -> bool:
+        return isinstance(
+            strategy,
+            (tf.distribute.TPUStrategy, tf.distribute.experimental.TPUStrategy),
+        )
+
+    def _has_sparsecore(self) -> bool:
+        strategy = tf.distribute.get_strategy()
+        if self._is_tpu_strategy(strategy):
+            tpu_embedding_feature = (
+                strategy.extended.tpu_hardware_feature.embedding_feature
+            )
+            return tpu_embedding_feature in (
+                EMBEDDING_FEATURE_V2,
+                EMBEDDING_FEATURE_V1,
+            )
+        return False
+
+    @keras_utils.no_automatic_dependency_tracking
+    def _sparsecore_init(
+        self,
+        feature_configs: dict[
+            str,
+            FeatureConfig | tf.tpu.experimental.embedding.FeatureConfig,
+        ],
+        table_stacking: str | Sequence[str] | Sequence[Sequence[str]],
+    ) -> None:
+        self._table_stacking = table_stacking
+
+        strategy = tf.distribute.get_strategy()
+        if not self._is_tpu_strategy(strategy):
+            raise ValueError(
+                "Placement to sparsecore was requested, however, we are not "
+                "running under a TPU strategy."
+            )
+
+        tpu_embedding_feature = (
+            strategy.extended.tpu_hardware_feature.embedding_feature
+        )
+
+        self._using_keras_rs_configuration = isinstance(
+            next(iter(feature_configs.values())), FeatureConfig
+        )
+        if self._using_keras_rs_configuration:
+            if self._sparse_core_embedding_config is not None:
+                raise ValueError(
+                    "The `sparse_core_embedding_config` argument is only "
+                    "supported when using "
+                    "`tf.tpu.experimental.embedding.FeatureConfig` instances "
+                    "for the configuration."
+                )
+            self._tpu_feature_configs, self._sparse_core_embedding_config = (
+                config_conversion.translate_keras_rs_configuration(
+                    feature_configs, table_stacking
+                )
+            )
+            if tpu_embedding_feature == EMBEDDING_FEATURE_V1:
+                # Remove auto-generated SparseCoreEmbeddingConfig, which is not
+                # used.
+                self._sparse_core_embedding_config = None
+        else:
+            if table_stacking != "auto":
+                raise ValueError(
+                    "The `table_stacking` argument is not supported when using "
+                    "`tf.tpu.experimental.embedding.FeatureConfig` for the "
+                    "configuration. You can use the `disable_table_stacking` "
+                    "attribute of "
+                    "`tf.tpu.experimental.embedding.SparseCoreEmbeddingConfig` "
+                    "to disable table stacking."
+                )
+            if (
+                tpu_embedding_feature == EMBEDDING_FEATURE_V1
+                and self._sparse_core_embedding_config is not None
+            ):
+                raise ValueError(
+                    "The `sparse_core_embedding_config` argument is not "
+                    "supported with this TPU generation."
+                )
+            self._tpu_feature_configs = (
+                config_conversion.clone_tf_feature_configs(feature_configs)
+            )
+
+        self._tpu_optimizer = config_conversion.translate_optimizer(
+            self._optimizer
+        )
+
+        if tpu_embedding_feature == EMBEDDING_FEATURE_V1:
+            self._tpu_embedding = tf.tpu.experimental.embedding.TPUEmbedding(
+                self._tpu_feature_configs,
+                self._tpu_optimizer,
+                self._pipeline_execution_with_tensor_core,
+            )
+            self._v1_call_id = 0
+        elif tpu_embedding_feature == EMBEDDING_FEATURE_V2:
+            self._tpu_embedding = tf.tpu.experimental.embedding.TPUEmbeddingV2(
+                self._tpu_feature_configs,
+                self._tpu_optimizer,
+                self._pipeline_execution_with_tensor_core,
+                self._sparse_core_embedding_config,
+            )
+        elif tpu_embedding_feature == UNSUPPORTED:
+            raise ValueError(
+                "Placement to sparsecore was requested, however, this TPU does "
+                "not support it."
+            )
+        elif tpu_embedding_feature != UNSUPPORTED:
+            raise ValueError(
+                f"Unsupported TPU embedding feature: {tpu_embedding_feature}."
+            )
+
+        # We need at least one trainable variable for the gradient trap to work.
+        # Note that the Python attribute name "_gradient_trap_dummy" should
+        # match the name of the variable GRADIENT_TRAP_DUMMY_NAME.
+        self._gradient_trap_dummy = self.add_weight(
+            name=GRADIENT_TRAP_DUMMY_NAME,
+            shape=(1,),
+            initializer=tf.zeros_initializer(),
+            trainable=True,
+            dtype=tf.float32,
+        )
+
+    def compute_output_shape(
+        self, input_shapes: types.Nested[types.Shape]
+    ) -> types.Nested[types.Shape]:
+        if self._using_keras_rs_configuration:
+            return super().compute_output_shape(input_shapes)
+
+        def _compute_output_shape(
+            feature_config: tf.tpu.experimental.embedding.FeatureConfig,
+            input_shape: types.Shape,
+        ) -> types.Shape:
+            if len(input_shape) < 1:
+                raise ValueError(
+                    f"Received input shape {input_shape}. Rank must be 1 or "
+                    "above."
+                )
+            max_sequence_length: int = feature_config.max_sequence_length
+            embed_dim = feature_config.table.dim
+            if (
+                feature_config.output_shape is not None
+                and feature_config.output_shape.rank is not None
+            ):
+                return tuple(feature_config.output_shape.as_list())
+            elif (
+                len(input_shape) == 2
+                and input_shape[-1] != 1
+                and max_sequence_length > 0
+            ):
+                # Update the input shape with the max sequence length. Only
+                # update when:
+                # 1. Input feature is 2D ragged or sparse tensor.
+                # 2. Output shape is not set and max sequence length is set.
+                return tuple(input_shape[:-1]) + (
+                    max_sequence_length,
+                    embed_dim,
+                )
+            elif len(input_shape) == 1:
+                return tuple(input_shape) + (embed_dim,)
+            else:
+                return tuple(input_shape[:-1]) + (embed_dim,)
+
+        output_shapes: types.Nested[types.Shape] = (
+            keras.tree.map_structure_up_to(
+                self._feature_configs,
+                _compute_output_shape,
+                self._feature_configs,
+                input_shapes,
+            )
+        )
+        return output_shapes
+
+    def _sparsecore_build(self, input_shapes: dict[str, types.Shape]) -> None:
+        if isinstance(
+            self._tpu_embedding, tf.tpu.experimental.embedding.TPUEmbedding
+        ):
+            tf_input_shapes = keras.tree.map_shape_structure(
+                tf.TensorShape, input_shapes
+            )
+            tpu_embedding_build = tf.autograph.to_graph(
+                self._tpu_embedding.build, recursive=False
+            )
+            tpu_embedding_build(
+                self._tpu_embedding, per_replica_input_shapes=tf_input_shapes
+            )
+        elif isinstance(
+            self._tpu_embedding, tf.tpu.experimental.embedding.TPUEmbeddingV2
+        ):
+            self._tpu_embedding.build()
+
+    def _sparsecore_call(
+        self,
+        inputs: dict[str, types.Tensor],
+        weights: dict[str, types.Tensor] | None = None,
+        training: bool = False,
+    ) -> dict[str, types.Tensor]:
+        del training  # Unused.
+        strategy = tf.distribute.get_strategy()
+        if not self._is_tpu_strategy(strategy):
+            raise RuntimeError(
+                "DistributedEmbedding needs to be called under a TPUStrategy "
+                "for features placed on the embedding feature but is being "
+                f"called under strategy {strategy}. Please use `strategy.run` "
+                "when calling this layer."
+            )
+        if isinstance(
+            self._tpu_embedding, tf.tpu.experimental.embedding.TPUEmbedding
+        ):
+            return self._tpu_embedding_lookup_v1(
+                self._tpu_embedding, inputs, weights
+            )
+        elif isinstance(
+            self._tpu_embedding, tf.tpu.experimental.embedding.TPUEmbeddingV2
+        ):
+            return self._tpu_embedding_lookup_v2(
+                self._tpu_embedding, inputs, weights
+            )
+        else:
+            raise ValueError(
+                "DistributedEmbedding is receiving features to lookup on the "
+                "TPU embedding feature but no such feature was configured."
+            )
+
+    def _sparsecore_get_embedding_tables(self) -> dict[str, types.Tensor]:
+        tables: dict[str, types.Tensor] = {}
+        strategy = tf.distribute.get_strategy()
+        # 4 is the number of sparsecores per chip
+        num_shards = strategy.num_replicas_in_sync * 4
+
+        def populate_table(
+            feature_config: tf.tpu.experimental.embedding.FeatureConfig,
+        ) -> None:
+            table_name = feature_config.table.name
+            if table_name in tables:
+                return
+
+            embedding_dim = feature_config.table.dim
+            table = self._tpu_embedding.embedding_tables[table_name]
+
+            # This table has num_sparse_cores mod shards, so we need to slice,
+            # reconcat and reshape.
+            table_shards = [
+                shard.numpy()[:, :embedding_dim] for shard in table.values
+            ]
+            full_table = keras.ops.concatenate(table_shards, axis=0)
+            full_table = keras.ops.concatenate(
+                keras.ops.split(full_table, num_shards, axis=0), axis=1
+            )
+            full_table = keras.ops.reshape(full_table, [-1, embedding_dim])
+            tables[table_name] = full_table[
+                : feature_config.table.vocabulary_size, :
+            ]
+
+        keras.tree.map_structure(populate_table, self._tpu_feature_configs)
+        return tables
+
+    def _verify_input_shapes(
+        self, input_shapes: types.Nested[types.Shape]
+    ) -> None:
+        if self._using_keras_rs_configuration:
+            return super()._verify_input_shapes(input_shapes)
+        # `tf.tpu.experimental.embedding.FeatureConfig` does not provide any
+        # information about the input shape, so there is nothing to verify.
+
+    def _tpu_embedding_lookup_v1(
+        self,
+        tpu_embedding: tf.tpu.experimental.embedding.TPUEmbedding,
+        inputs: dict[str, types.Tensor],
+        weights: dict[str, types.Tensor] | None = None,
+    ) -> dict[str, types.Tensor]:
+        # Each call to this function increments the _v1_call_id by 1, this
+        # allows us to tag each of the main embedding ops with this call id so
+        # that we know during graph rewriting passes which ops correspond to the
+        # same layer call.
+        self._v1_call_id += 1
+        name = str(self._v1_call_id)
+
+        # Set training to true, even during eval. When name is set, this will
+        # trigger a pass that updates the training based on if there is a send
+        # gradients with the same name.
+        tpu_embedding.enqueue(inputs, weights, training=True, name=name)
+
+        @tf.custom_gradient  # type: ignore
+        def gradient_trap(
+            dummy: types.Tensor,
+        ) -> tuple[
+            list[types.Tensor], Callable[[tuple[types.Tensor]], types.Tensor]
+        ]:
+            """Register a gradient function for activation."""
+            activations = tpu_embedding.dequeue(name=name)
+
+            def grad(*grad_wrt_activations: types.Tensor) -> types.Tensor:
+                """Gradient function."""
+                # Since the output were flattened, the gradients are also
+                # flattened. Pack them back into the correct nested structure.
+                gradients = tf.nest.pack_sequence_as(
+                    self._placement_to_path_to_feature_config["sparsecore"],
+                    grad_wrt_activations,
+                )
+                tpu_embedding.apply_gradients(gradients, name=name)
+
+                # This is the gradient for the input variable.
+                return tf.zeros_like(dummy)
+
+            # Custom gradient functions don't like nested structures of tensors,
+            # so we flatten them here.
+            return tf.nest.flatten(activations), grad
+
+        activations_with_trap = gradient_trap(self._gradient_trap_dummy.value)
+        result: dict[str, types.Tensor] = tf.nest.pack_sequence_as(
+            self._placement_to_path_to_feature_config["sparsecore"],
+            activations_with_trap,
+        )
+        return result
+
+    def _tpu_embedding_lookup_v2(
+        self,
+        tpu_embedding: tf.tpu.experimental.embedding.TPUEmbeddingV2,
+        inputs: dict[str, types.Tensor],
+        weights: dict[str, types.Tensor] | None = None,
+    ) -> dict[str, types.Tensor]:
+        @tf.custom_gradient  # type: ignore
+        def gradient_trap(
+            dummy: types.Tensor,
+        ) -> tuple[
+            list[types.Tensor], Callable[[tuple[types.Tensor]], types.Tensor]
+        ]:
+            """Register a gradient function for activation."""
+            activations, preserved_result = tpu_embedding(inputs, weights)
+
+            def grad(*grad_wrt_activations: types.Tensor) -> types.Tensor:
+                """Gradient function."""
+                # Since the output were flattened, the gradients are also
+                # flattened. Pack them back into the correct nested structure.
+                gradients = tf.nest.pack_sequence_as(
+                    self._placement_to_path_to_feature_config["sparsecore"],
+                    grad_wrt_activations,
+                )
+                tpu_embedding.apply_gradients(
+                    gradients, preserved_outputs=preserved_result
+                )
+                # This is the gradient for the input variable.
+                return tf.zeros_like(dummy)
+
+            # Custom gradient functions don't like nested structures of tensors,
+            # so we flatten them here.
+            return tf.nest.flatten(activations), grad
+
+        activations_with_trap = gradient_trap(self._gradient_trap_dummy)
+        result: dict[str, types.Tensor] = tf.nest.pack_sequence_as(
+            self._placement_to_path_to_feature_config["sparsecore"],
+            activations_with_trap,
+        )
+        return result
+
+    def _trackable_children(
+        self, save_type: str = "checkpoint", **kwargs: dict[str, Any]
+    ) -> dict[str, Any]:
+        # Remove dummy variable, we don't want it in checkpoints.
+        children: dict[str, Any] = super()._trackable_children(
+            save_type, **kwargs
+        )
+        children.pop(GRADIENT_TRAP_DUMMY_NAME, None)
+        return children
+
+
+DistributedEmbedding.__doc__ = (
+    base_distributed_embedding.DistributedEmbedding.__doc__
+)

keras_rs/src/layers/feature_interaction/dot_interaction.py

@@ -205,8 +205,8 @@ class DotInteraction(keras.layers.Layer):
         return activations
 
     def compute_output_shape(
-        self, input_shape: list[types.TensorShape]
-    ) -> types.TensorShape:
+        self, input_shape: list[types.Shape]
+    ) -> types.Shape:
         num_features = len(input_shape)
         batch_size = input_shape[0][0]
 

keras_rs/src/layers/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
@@ -92,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)
@@ -129,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:
@@ -155,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,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Any
 
 import keras
 
@@ -55,8 +55,8 @@ class BruteForceRetrieval(Retrieval):
 
     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,
@@ -81,7 +81,7 @@ class BruteForceRetrieval(Retrieval):
     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.
 
@@ -125,7 +125,7 @@ class BruteForceRetrieval(Retrieval):
 
     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:

keras_rs/src/layers/retrieval/retrieval.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Optional, Union
+from typing import Any
 
 import keras
 
@@ -35,7 +35,7 @@ class Retrieval(keras.layers.Layer, abc.ABC):
     def _validate_candidate_embeddings_and_ids(
         self,
         candidate_embeddings: types.Tensor,
-        candidate_ids: Optional[types.Tensor] = None,
+        candidate_ids: types.Tensor | None = None,
     ) -> None:
         """Validates inputs to `update_candidates()`."""
 
@@ -71,7 +71,7 @@ class Retrieval(keras.layers.Layer, abc.ABC):
     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.
 
@@ -85,7 +85,7 @@ class Retrieval(keras.layers.Layer, abc.ABC):
     @abc.abstractmethod
     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:

keras_rs/src/losses/pairwise_loss.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Any, Optional
+from typing import Any
 
 import keras
 from keras import ops
@@ -43,7 +43,7 @@ class PairwiseLoss(keras.losses.Loss, abc.ABC):
         self,
         labels: types.Tensor,
         logits: types.Tensor,
-        mask: Optional[types.Tensor] = None,
+        mask: types.Tensor | None = None,
     ) -> tuple[types.Tensor, types.Tensor]:
         # Mask all values less than 0 (since less than 0 implies invalid
         # labels).

keras_rs/src/losses/pairwise_mean_squared_error.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from keras import ops
 
 from keras_rs.src import types
@@ -20,7 +18,7 @@ class PairwiseMeanSquaredError(PairwiseLoss):
         self,
         labels: types.Tensor,
         logits: types.Tensor,
-        mask: Optional[types.Tensor] = None,
+        mask: types.Tensor | None = None,
     ) -> tuple[types.Tensor, types.Tensor]:
         # Override `PairwiseLoss.compute_unreduced_loss` since pairwise weights
         # for MSE are computed differently.

keras_rs/src/metrics/dcg.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 from keras import ops
 from keras.saving import deserialize_keras_object
@@ -25,7 +25,7 @@ from keras_rs.src.utils.doc_string_utils import format_docstring
 class DCG(RankingMetric):
     def __init__(
         self,
-        k: Optional[int] = None,
+        k: int | None = None,
         gain_fn: Callable[[types.Tensor], types.Tensor] = default_gain_fn,
         rank_discount_fn: Callable[
             [types.Tensor], types.Tensor

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],

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

keras_rs/src/metrics/ndcg.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional
+from typing import Any, Callable
 
 from keras import ops
 from keras.saving import deserialize_keras_object
@@ -25,7 +25,7 @@ from keras_rs.src.utils.doc_string_utils import format_docstring
 class NDCG(RankingMetric):
     def __init__(
         self,
-        k: Optional[int] = None,
+        k: int | None = None,
         gain_fn: Callable[[types.Tensor], types.Tensor] = default_gain_fn,
         rank_discount_fn: Callable[
             [types.Tensor], types.Tensor

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