Updated code

Author: LakshmiKalaKadali

keras_rs/src/losses/list_mle_loss.py

@@ -127,17 +127,19 @@ def compute_unreduced_loss(
         )
         sorted_logits = ops.subtract(sorted_logits, raw_max)
 
-        exp_logits = ops.exp(sorted_logits)
-        exp_logits = ops.where(
-            sorted_valid_mask, exp_logits, ops.zeros_like(exp_logits)
+        # 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)

keras_rs/src/losses/list_mle_loss_test.py

@@ -10,16 +10,16 @@
 
 class ListMLELossTest(testing.TestCase, parameterized.TestCase):
     def setUp(self):
-        self.unbatched_scores = ops.array([1.0, 3.0, 2.0, 4.0, 0.8])
-        self.unbatched_labels = ops.array([1.0, 0.0, 1.0, 3.0, 2.0])
+        self.unbatched_scores = ops.array([1.0, 3.0, 2.0, 4.0, 0.8], dtype="float32")
+        self.unbatched_labels = ops.array([1.0, 0.0, 1.0, 3.0, 2.0], dtype="float32")
 
         self.batched_scores = ops.array(
-            [[1.0, 3.0, 2.0, 4.0, 0.8], [1.0, 1.8, 2.0, 3.0, 2.0]]
+            [[1.0, 3.0, 2.0, 4.0, 0.8], [1.0, 1.8, 2.0, 3.0, 2.0]], dtype="float32"
         )
         self.batched_labels = ops.array(
-            [[1.0, 0.0, 1.0, 3.0, 2.0], [0.0, 1.0, 2.0, 3.0, 1.5]]
+            [[1.0, 0.0, 1.0, 3.0, 2.0], [0.0, 1.0, 2.0, 3.0, 1.5]], dtype="float32"
         )
-        self.expected_output = ops.array([6.865693, 3.088192])
+        self.expected_output = ops.array([6.865693, 3.088192], dtype="float32")
 
     def test_unbatched_input(self):
         loss = ListMLELoss(reduction="none")
@@ -43,7 +43,6 @@ def test_temperature(self):
         output_temp = loss_temp(
             y_true=self.batched_labels, y_pred=self.batched_scores
         )
-
         self.assertAllClose(
             output_temp,
             [10.969891, 2.1283305],
@@ -60,7 +59,6 @@ def test_invalid_input_rank(self):
     def test_loss_reduction(self):
         loss = ListMLELoss(reduction="sum_over_batch_size")
         output = loss(y_true=self.batched_labels, y_pred=self.batched_scores)
-
         self.assertAlmostEqual(
             ops.convert_to_numpy(output), 4.9769425, places=5
         )