keras-rs-nightly 0.3.1.dev202510100326__tar.gz → 0.3.1.dev202511120334__tar.gz

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202510100326
+Version: 0.3.1.dev202511120334
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/api/losses/__init__.py

@@ -4,6 +4,7 @@ This file was autogenerated. Do not edit it by hand,
 since your modifications would be overwritten.
 """
 
+from keras_rs.src.losses.list_mle_loss import ListMLELoss as ListMLELoss
 from keras_rs.src.losses.pairwise_hinge_loss import (
     PairwiseHingeLoss as PairwiseHingeLoss,
 )

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/src/layers/embedding/jax/distributed_embedding.py

@@ -9,6 +9,7 @@ import keras
 import numpy as np
 from jax import numpy as jnp
 from jax.experimental import layout as jax_layout
+from jax.experimental import multihost_utils
 from jax_tpu_embedding.sparsecore.lib.nn import embedding
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 from jax_tpu_embedding.sparsecore.lib.nn import (
@@ -442,7 +443,50 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
 
         # Collect all stacked tables.
         table_specs = embedding.get_table_specs(feature_specs)
-        table_stacks = embedding_utils.get_table_stacks(table_specs)
+        table_stacks = jte_table_stacking.get_table_stacks(table_specs)
+
+        # Update stacked table stats to max of values across involved tables.
+        max_ids_per_partition = {}
+        max_unique_ids_per_partition = {}
+        required_buffer_size_per_device = {}
+        id_drop_counters = {}
+        for stack_name, stack in table_stacks.items():
+            max_ids_per_partition[stack_name] = np.max(
+                np.asarray(
+                    [s.max_ids_per_partition for s in stack], dtype=np.int32
+                )
+            )
+            max_unique_ids_per_partition[stack_name] = np.max(
+                np.asarray(
+                    [s.max_unique_ids_per_partition for s in stack],
+                    dtype=np.int32,
+                )
+            )
+
+            # Only set the suggested buffer size if set on any individual table.
+            valid_buffer_sizes = [
+                s.suggested_coo_buffer_size_per_device
+                for s in stack
+                if s.suggested_coo_buffer_size_per_device is not None
+            ]
+            if valid_buffer_sizes:
+                required_buffer_size_per_device[stack_name] = np.max(
+                    np.asarray(valid_buffer_sizes, dtype=np.int32)
+                )
+
+            id_drop_counters[stack_name] = 0
+
+        aggregated_stats = embedding.SparseDenseMatmulInputStats(
+            max_ids_per_partition=max_ids_per_partition,
+            max_unique_ids_per_partition=max_unique_ids_per_partition,
+            required_buffer_size_per_sc=required_buffer_size_per_device,
+            id_drop_counters=id_drop_counters,
+        )
+        embedding.update_preprocessing_parameters(
+            feature_specs,
+            aggregated_stats,
+            num_sc_per_device,
+        )
 
         # Create variables for all stacked tables and slot variables.
         with sparsecore_distribution.scope():
@@ -516,7 +560,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
 
         # Each stacked-table gets a ShardedCooMatrix.
         table_specs = embedding.get_table_specs(self._config.feature_specs)
-        table_stacks = embedding_utils.get_table_stacks(table_specs)
+        table_stacks = jte_table_stacking.get_table_stacks(table_specs)
         stacked_table_specs = {
             stack_name: stack[0].stacked_table_spec
             for stack_name, stack in table_stacks.items()
@@ -600,31 +644,26 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             # underlying stacked tables specs in the feature specs.
 
             # Aggregate stats across all processes/devices via pmax.
-            num_local_cpu_devices = jax.local_device_count("cpu")
-
-            def pmax_aggregate(x: Any) -> Any:
-                if not hasattr(x, "ndim"):
-                    x = np.array(x)
-                tiled_x = np.tile(x, (num_local_cpu_devices, *([1] * x.ndim)))
-                return jax.pmap(
-                    lambda y: jax.lax.pmax(y, "all_cpus"),  # type: ignore[no-untyped-call]
-                    axis_name="all_cpus",
-                    backend="cpu",
-                )(tiled_x)[0]
-
-            full_stats = jax.tree.map(pmax_aggregate, stats)
+            all_stats = multihost_utils.process_allgather(stats)
+            aggregated_stats = jax.tree.map(
+                lambda x: jnp.max(x, axis=0), all_stats
+            )
 
             # Check if stats changed enough to warrant action.
             stacked_table_specs = embedding.get_stacked_table_specs(
                 self._config.feature_specs
             )
             changed = any(
-                np.max(full_stats.max_ids_per_partition[stack_name])
+                np.max(aggregated_stats.max_ids_per_partition[stack_name])
                 > spec.max_ids_per_partition
-                or np.max(full_stats.max_unique_ids_per_partition[stack_name])
+                or np.max(
+                    aggregated_stats.max_unique_ids_per_partition[stack_name]
+                )
                 > spec.max_unique_ids_per_partition
                 or (
-                    np.max(full_stats.required_buffer_size_per_sc[stack_name])
+                    np.max(
+                        aggregated_stats.required_buffer_size_per_sc[stack_name]
+                    )
                     * num_sc_per_device
                 )
                 > (spec.suggested_coo_buffer_size_per_device or 0)
@@ -634,7 +673,9 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
             # Update configuration and repeat preprocessing if stats changed.
             if changed:
                 embedding.update_preprocessing_parameters(
-                    self._config.feature_specs, full_stats, num_sc_per_device
+                    self._config.feature_specs,
+                    aggregated_stats,
+                    num_sc_per_device,
                 )
 
                 # Re-execute preprocessing with consistent input statistics.
@@ -720,7 +761,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
         config = self._config
         num_table_shards = config.mesh.devices.size * config.num_sc_per_device
         table_specs = embedding.get_table_specs(config.feature_specs)
-        sharded_tables = embedding_utils.stack_and_shard_tables(
+        sharded_tables = jte_table_stacking.stack_and_shard_tables(
             table_specs,
             tables,
             num_table_shards,
@@ -763,7 +804,7 @@ class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
 
         return typing.cast(
             dict[str, ArrayLike],
-            embedding_utils.unshard_and_unstack_tables(
+            jte_table_stacking.unshard_and_unstack_tables(
                 table_specs, table_variables, num_table_shards
             ),
         )

keras_rs_nightly-0.3.1.dev202511120334/keras_rs/src/layers/embedding/jax/embedding_utils.py

@@ -0,0 +1,244 @@
+"""Utility functions for manipulating JAX embedding tables and inputs."""
+
+import collections
+from typing import Any, Mapping, NamedTuple, Sequence, TypeAlias, TypeVar
+
+import jax
+import numpy as np
+from jax_tpu_embedding.sparsecore.lib.nn import embedding
+from jax_tpu_embedding.sparsecore.lib.nn import table_stacking
+from jax_tpu_embedding.sparsecore.lib.nn.embedding_spec import FeatureSpec
+
+from keras_rs.src.types import Nested
+
+T = TypeVar("T")
+
+# Any to support tf.Ragged without needing an explicit TF dependency.
+ArrayLike: TypeAlias = jax.Array | np.ndarray | Any  # type: ignore
+Shape: TypeAlias = tuple[int, ...]
+
+
+class FeatureSamples(NamedTuple):
+    tokens: ArrayLike
+    weights: ArrayLike
+
+
+class ShardedCooMatrix(NamedTuple):
+    shard_starts: ArrayLike
+    shard_ends: ArrayLike
+    col_ids: ArrayLike
+    row_ids: ArrayLike
+    values: ArrayLike
+
+
+def convert_to_numpy(
+    ragged_or_dense: np.ndarray[Any, Any] | Sequence[Sequence[Any]] | Any,
+    dtype: Any,
+) -> np.ndarray[Any, Any]:
+    """Converts a ragged or dense list of inputs to a ragged/dense numpy array.
+
+    The output is adjusted to be 2D.
+
+    Args:
+        ragged_or_dense: Input that is either already a numpy array, or nested
+            sequence.
+        dtype: Numpy dtype of output array.
+
+    Returns:
+        Corresponding numpy array.
+    """
+    if hasattr(ragged_or_dense, "numpy"):
+        # Support tf.RaggedTensor and other TF input dtypes.
+        if callable(getattr(ragged_or_dense, "numpy")):
+            ragged_or_dense = ragged_or_dense.numpy()
+
+    if isinstance(ragged_or_dense, jax.Array):
+        ragged_or_dense = np.asarray(ragged_or_dense)
+
+    if isinstance(ragged_or_dense, np.ndarray):
+        # Convert 1D to 2D.
+        if ragged_or_dense.dtype != np.ndarray and ragged_or_dense.ndim == 1:
+            return ragged_or_dense.reshape(-1, 1).astype(dtype)
+
+        # If dense, return converted dense type.
+        if ragged_or_dense.dtype != np.ndarray:
+            return ragged_or_dense.astype(dtype)
+
+        # Ragged numpy array.
+        return ragged_or_dense
+
+    # Handle 1D sequence input.
+    if not isinstance(ragged_or_dense[0], collections.abc.Sequence):
+        return np.asarray(ragged_or_dense, dtype=dtype).reshape(-1, 1)
+
+    # Assemble elements into an nd-array.
+    counts = [len(vals) for vals in ragged_or_dense]
+    if all([count == counts[0] for count in counts]):
+        # Dense input.
+        return np.asarray(ragged_or_dense, dtype=dtype)
+    else:
+        # Ragged input, convert to ragged numpy arrays.
+        return np.array(
+            [np.array(row, dtype=dtype) for row in ragged_or_dense],
+            dtype=np.ndarray,
+        )
+
+
+def ones_like(
+    ragged_or_dense: np.ndarray[Any, Any], dtype: Any = None
+) -> np.ndarray[Any, Any]:
+    """Creates an array of ones the same as as the input.
+
+    This differs from traditional numpy in that a ragged input will lead to
+    a resulting ragged array of ones, whereas np.ones_like(...) will instead
+    only consider the outer array and return a 1D dense array of ones.
+
+    Args:
+        ragged_or_dense: The ragged or dense input whose shape and data-type
+            define these same attributes of the returned array.
+        dtype: The data-type of the returned array.
+
+    Returns:
+        An array of ones with the same shape as the input, and specified data
+        type.
+    """
+    dtype = dtype or ragged_or_dense.dtype
+    if ragged_or_dense.dtype == np.ndarray:
+        # Ragged.
+        return np.array(
+            [np.ones_like(row, dtype=dtype) for row in ragged_or_dense],
+            dtype=np.ndarray,
+        )
+    else:
+        # Dense.
+        return np.ones_like(ragged_or_dense, dtype=dtype)
+
+
+def create_feature_samples(
+    feature_structure: Nested[T],
+    feature_ids: Nested[ArrayLike | Sequence[int] | Sequence[Sequence[int]]],
+    feature_weights: None
+    | (Nested[ArrayLike | Sequence[float] | Sequence[Sequence[float]]]),
+) -> Nested[FeatureSamples]:
+    """Constructs a collection of sample tuples from provided IDs and weights.
+
+    Args:
+        feature_structure: The nested structure of the inputs (typically
+            `FeatureSpec`s).
+        feature_ids: The feature IDs to use for the samples.
+        feature_weights: The feature weights to use for the samples.  Defaults
+            to ones if not provided.
+
+    Returns:
+        A nested collection of `FeatureSamples` corresponding to the input IDs
+        and weights, for use in embedding lookups.
+    """
+    # Create numpy arrays from inputs.
+    feature_ids = jax.tree.map(
+        lambda _, ids: convert_to_numpy(ids, np.int32),
+        feature_structure,
+        feature_ids,
+    )
+
+    if feature_weights is None:
+        # Make ragged or dense ones_like.
+        feature_weights = jax.tree.map(
+            lambda _, ids: ones_like(ids, np.float32),
+            feature_structure,
+            feature_ids,
+        )
+    else:
+        feature_weights = jax.tree.map(
+            lambda _, wgts: convert_to_numpy(wgts, np.float32),
+            feature_structure,
+            feature_weights,
+        )
+
+    # Assemble.
+    def _create_feature_samples(
+        sample_ids: np.ndarray[Any, Any],
+        sample_weights: np.ndarray[Any, Any],
+    ) -> FeatureSamples:
+        return FeatureSamples(sample_ids, sample_weights)
+
+    output: Nested[FeatureSamples] = jax.tree.map(
+        lambda _, sample_ids, sample_weights: _create_feature_samples(
+            sample_ids, sample_weights
+        ),
+        feature_structure,
+        feature_ids,
+        feature_weights,
+    )
+    return output
+
+
+def stack_and_shard_samples(
+    feature_specs: Nested[FeatureSpec],
+    feature_samples: Nested[FeatureSamples],
+    local_device_count: int,
+    global_device_count: int,
+    num_sc_per_device: int,
+    static_buffer_size: int | Mapping[str, int] | None = None,
+) -> tuple[dict[str, ShardedCooMatrix], embedding.SparseDenseMatmulInputStats]:
+    """Prepares input samples for use in embedding lookups.
+
+    Args:
+        feature_specs: Nested collection of feature specifications.
+        feature_samples: Nested collection of feature samples.
+        local_device_count: Number of local JAX devices.
+        global_device_count: Number of global JAX devices.
+        num_sc_per_device: Number of sparsecores per device.
+        static_buffer_size: The static buffer size to use for the samples.
+          Defaults to None, in which case an upper-bound for the buffer size
+          will be automatically determined.
+
+    Returns:
+        The preprocessed inputs, and statistics useful for updating FeatureSpecs
+        based on the provided input data.
+    """
+    del static_buffer_size  # Currently ignored.
+    flat_feature_specs, _ = jax.tree.flatten(feature_specs)
+
+    feature_tokens = []
+    feature_weights = []
+
+    def collect_tokens_and_weights(
+        feature_spec: FeatureSpec, samples: FeatureSamples
+    ) -> None:
+        del feature_spec
+        feature_tokens.append(samples.tokens)
+        feature_weights.append(samples.weights)
+
+    jax.tree.map(collect_tokens_and_weights, feature_specs, feature_samples)
+
+    preprocessed_inputs, stats = embedding.preprocess_sparse_dense_matmul_input(
+        feature_tokens,
+        feature_weights,
+        flat_feature_specs,
+        local_device_count=local_device_count,
+        global_device_count=global_device_count,
+        num_sc_per_device=num_sc_per_device,
+        sharding_strategy="MOD",
+        has_leading_dimension=False,
+        allow_id_dropping=True,
+    )
+
+    out: dict[str, ShardedCooMatrix] = {}
+    tables_names = preprocessed_inputs.lhs_row_pointers.keys()
+    for table_name in tables_names:
+        shard_ends = preprocessed_inputs.lhs_row_pointers[table_name]
+        shard_starts = np.concatenate(
+            [
+                np.asarray([0]),
+                table_stacking._next_largest_multiple(shard_ends[:-1], 8),
+            ]
+        )
+        out[table_name] = ShardedCooMatrix(
+            shard_starts=shard_starts,
+            shard_ends=shard_ends,
+            col_ids=preprocessed_inputs.lhs_embedding_ids[table_name],
+            row_ids=preprocessed_inputs.lhs_sample_ids[table_name],
+            values=preprocessed_inputs.lhs_gains[table_name],
+        )
+
+    return out, stats

keras_rs_nightly-0.3.1.dev202511120334/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_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/src/metrics/ranking_metrics_utils.py

@@ -85,6 +85,25 @@ def sort_by_scores(
     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:

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs/src/version.py

@@ -1,7 +1,7 @@
 from keras_rs.src.api_export import keras_rs_export
 
 # Unique source of truth for the version number.
-__version__ = "0.3.1.dev202510100326"
+__version__ = "0.3.1.dev202511120334"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.3.1.dev202510100326
+Version: 0.3.1.dev202511120334
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_rs_nightly-0.3.1.dev202510100326 → keras_rs_nightly-0.3.1.dev202511120334}/keras_rs_nightly.egg-info/SOURCES.txt

@@ -33,6 +33,7 @@ keras_rs/src/layers/retrieval/remove_accidental_hits.py
 keras_rs/src/layers/retrieval/retrieval.py
 keras_rs/src/layers/retrieval/sampling_probability_correction.py
 keras_rs/src/losses/__init__.py
+keras_rs/src/losses/list_mle_loss.py
 keras_rs/src/losses/pairwise_hinge_loss.py
 keras_rs/src/losses/pairwise_logistic_loss.py
 keras_rs/src/losses/pairwise_loss.py