Used sort_by_score function instead of manual sorting

Author: LakshmiKalaKadali

api_gen.py

@@ -8,7 +8,6 @@
 
 import os
 import shutil
-import pre_commit
 import namex
 
 PACKAGE = "keras_rs"

keras_rs/api/layers/__init__.py

@@ -4,15 +4,34 @@
 since your modifications would be overwritten.
 """
 
-
-from keras_rs.src.layers.embedding.distributed_embedding import DistributedEmbedding as DistributedEmbedding
-from keras_rs.src.layers.embedding.distributed_embedding_config import FeatureConfig as FeatureConfig
-from keras_rs.src.layers.embedding.distributed_embedding_config import TableConfig as TableConfig
-from keras_rs.src.layers.embedding.embed_reduce import EmbedReduce as EmbedReduce
-from keras_rs.src.layers.feature_interaction.dot_interaction import DotInteraction as DotInteraction
-from keras_rs.src.layers.feature_interaction.feature_cross import FeatureCross as FeatureCross
-from keras_rs.src.layers.retrieval.brute_force_retrieval import BruteForceRetrieval as BruteForceRetrieval
-from keras_rs.src.layers.retrieval.hard_negative_mining import HardNegativeMining as HardNegativeMining
-from keras_rs.src.layers.retrieval.remove_accidental_hits import RemoveAccidentalHits as RemoveAccidentalHits
+from keras_rs.src.layers.embedding.distributed_embedding import (
+    DistributedEmbedding as DistributedEmbedding,
+)
+from keras_rs.src.layers.embedding.distributed_embedding_config import (
+    FeatureConfig as FeatureConfig,
+)
+from keras_rs.src.layers.embedding.distributed_embedding_config import (
+    TableConfig as TableConfig,
+)
+from keras_rs.src.layers.embedding.embed_reduce import (
+    EmbedReduce as EmbedReduce,
+)
+from keras_rs.src.layers.feature_interaction.dot_interaction import (
+    DotInteraction as DotInteraction,
+)
+from keras_rs.src.layers.feature_interaction.feature_cross import (
+    FeatureCross as FeatureCross,
+)
+from keras_rs.src.layers.retrieval.brute_force_retrieval import (
+    BruteForceRetrieval as BruteForceRetrieval,
+)
+from keras_rs.src.layers.retrieval.hard_negative_mining import (
+    HardNegativeMining as HardNegativeMining,
+)
+from keras_rs.src.layers.retrieval.remove_accidental_hits import (
+    RemoveAccidentalHits as RemoveAccidentalHits,
+)
 from keras_rs.src.layers.retrieval.retrieval import Retrieval as Retrieval
-from keras_rs.src.layers.retrieval.sampling_probability_correction import SamplingProbabilityCorrection as SamplingProbabilityCorrection
+from keras_rs.src.layers.retrieval.sampling_probability_correction import (
+    SamplingProbabilityCorrection as SamplingProbabilityCorrection,
+)

keras_rs/src/losses/list_mle_loss.py

@@ -6,6 +6,8 @@
 from keras_rs.src import types
 from keras_rs.src.metrics.utils import standardize_call_inputs_ranks
 from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+
 
 @keras_rs_export("keras_rs.losses.ListMLELoss")
 class ListMLELoss(keras.losses.Loss):
@@ -20,7 +22,7 @@ class ListMLELoss(keras.losses.Loss):
 
     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)))
     ```
@@ -65,7 +67,6 @@ def __init__(self, temperature: float = 1.0, **kwargs: Any) -> None:
         self.temperature = temperature
         self._epsilon = 1e-10  
 
-
     def compute_unreduced_loss(
         self,
         labels: types.Tensor,
@@ -84,7 +85,7 @@ def compute_unreduced_loss(
             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:
@@ -95,13 +96,17 @@ def compute_unreduced_loss(
 
         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_indices = ops.argsort(-labels_for_sorting, axis=-1)
-        
-        sorted_logits = ops.take_along_axis(logits_masked, sorted_indices, axis=-1)
-        sorted_valid_mask = ops.take_along_axis(valid_mask, sorted_indices, axis=-1)
+    
+        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,
@@ -111,26 +116,21 @@ def compute_unreduced_loss(
         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 = sorted_logits - raw_max
-    
-        
+
         exp_logits = ops.exp(sorted_logits)
-       
         exp_logits = ops.where(sorted_valid_mask, exp_logits, ops.zeros_like(exp_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 = sorted_logits - log_normalizers
 
         log_probs = ops.where(sorted_valid_mask, log_probs, ops.zeros_like(log_probs))
 
-
         negative_log_likelihood = -ops.sum(log_probs, axis=1, keepdims=True)
 
         negative_log_likelihood = ops.where(batch_has_valid_items, negative_log_likelihood, 
@@ -187,9 +187,8 @@ def call(
         losses = ops.multiply(losses, weights)
         losses = ops.squeeze(losses, axis=-1)  
         return losses
-          
 
     def get_config(self) -> dict[str, Any]:
         config: dict[str, Any] = super().get_config()
         config.update({"temperature": self.temperature})
-        return config
+        return config