Updated code with 'cumsum' function

Author: LakshmiKalaKadali

keras_rs/src/losses/list_mle_loss.py

@@ -136,9 +136,12 @@ def compute_unreduced_loss(
         # 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)
-        cumsum_forward = ops.cumsum(exp_logits, axis=1)
-        total_sum = ops.sum(exp_logits, axis=1, keepdims=True)
-        cumsum_from_right = total_sum - cumsum_forward + exp_logits
+        #cumsum_forward = ops.cumsum(exp_logits, axis=1)
+        #total_sum = ops.sum(exp_logits, axis=1, keepdims=True)
+        #cumsum_from_right = total_sum - cumsum_forward + 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 = ops.subtract(sorted_logits, log_normalizers)

keras_rs/src/losses/list_mle_loss_test.py

@@ -16,7 +16,6 @@ def setUp(self):
         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]],
             dtype="float32",