Modified code with axis

Author: LakshmiKalaKadali

keras_rs/src/losses/list_mle_loss.py

@@ -136,12 +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
-        reversed_exp = ops.flip(exp_logits, 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
+        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_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)
@@ -211,6 +211,7 @@ def call(
         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()

keras_rs/src/losses/list_mle_loss_test.py

@@ -7,7 +7,6 @@
 from keras_rs.src import testing
 from keras_rs.src.losses.list_mle_loss import ListMLELoss
 
-
 class ListMLELossTest(testing.TestCase, parameterized.TestCase):
     def setUp(self):
         self.unbatched_scores = ops.array(