Handled ops.flip for torch backend

Author: LakshmiKalaKadali

keras_rs/src/losses/list_mle_loss.py

@@ -129,17 +129,19 @@ def compute_unreduced_loss(
 
         # Set invalid positions to very negative BEFORE exp
         sorted_logits = ops.where(
-        sorted_valid_mask, sorted_logits, ops.full_like(sorted_logits, -1e9)
+            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)
+        # 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
 
         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,14 +10,20 @@
 
 class ListMLELossTest(testing.TestCase, parameterized.TestCase):
     def setUp(self):
-        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.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]], dtype="float32"
+            [[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]], dtype="float32"
+            [[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], dtype="float32")