axis values unchanged

Author: LakshmiKalaKadali

keras_rs/src/losses/list_mle_loss.py

@@ -105,15 +105,13 @@ def compute_unreduced_loss(
         logits_masked = ops.where(
             valid_mask, logits, ops.full_like(logits, -1e9)
         )
-
         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, ops.cast(self.temperature, dtype=sorted_logits.dtype)
         )
@@ -139,9 +137,9 @@ def compute_unreduced_loss(
         # 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_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)

keras_rs/src/losses/list_mle_loss_test.py

@@ -7,6 +7,7 @@
 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(