revert optimization from PR#4166

Author: NuojCheng

src/maxtext/kernels/ragged/ragged_sort.py

@@ -186,6 +186,8 @@ def _ring_ragged_sort_bwd(res, g_out):
           valid_rows_mask=valid_rows_mask,
           reduce_group_size=topk,
           enforce_fallback=enforce_gather_reduce_fallback,
+          flops_override=gather_reduce_flops_override,
+          bytes_accessed_override=gather_reduce_bytes_accessed_override,
       )
     else:
       # Buffering: g_x has size `local_buffer_size` (packed).
@@ -209,6 +211,8 @@ def _ring_ragged_sort_bwd(res, g_out):
           valid_rows_mask=valid_rows_mask,
           reduce_group_size=topk,
           enforce_fallback=enforce_gather_reduce_fallback,
+          flops_override=gather_reduce_flops_override,
+          bytes_accessed_override=gather_reduce_bytes_accessed_override,
       )
     return grad_hidden_states, None
 

src/maxtext/layers/moe.py

@@ -918,6 +918,7 @@ def permute(
           repeats=group_size,
           total_repeat_length=math.prod(selected_experts.shape),
       )
+
     return (
         sorted_inputs,
         sorted_selected_experts,