Fixing kernel to accomodate compiler change.

PiperOrigin-RevId: 895535377

Author: clee1994

src/maxtext/kernels/gather_reduce_sc.py

@@ -460,14 +460,14 @@ def perform_add(
           weights_evens = tpu.unpack_subelements(
               _F32[16],
               raw_weights,
-              0,
+              1,
               ir.Attribute.parse("#tpu.pack_format<interleaved>"),
           )
           # part=0 corresponds to sc_tpu.unpackf result 0 (Odd indices)
           weights_odds = tpu.unpack_subelements(
               _F32[16],
               raw_weights,
-              1,
+              0,
               ir.Attribute.parse("#tpu.pack_format<interleaved>"),
           )
 
@@ -541,13 +541,13 @@ def get_row_val(row_idx):
               vec_f32_evens = tpu.unpack_subelements(
                   _F32[16],
                   vec_bf16_2x16,
-                  0,
+                  1,
                   ir.Attribute.parse("#tpu.pack_format<interleaved>"),
               )
               vec_f32_odds = tpu.unpack_subelements(
                   _F32[16],
                   vec_bf16_2x16,
-                  1,
+                  0,
                   ir.Attribute.parse("#tpu.pack_format<interleaved>"),
               )
               parity_of_row = vector.extract(
@@ -620,7 +620,7 @@ def get_row_val(row_idx):
 
           packed = tpu.pack_subelements(
               _BF16[2, vreg_size],
-              [row8, row0],
+              [row0, row8],
               [0, 1],
               ir.Attribute.parse("#tpu.pack_format<interleaved>"),
           )