Merge pull request #4080 from AI-Hypercomputer:chengnuojin-rs-g3

PiperOrigin-RevId: 927429580

Author: Google-ML-Automation

src/maxtext/kernels/ragged/ragged_gather_reduce.py

@@ -29,9 +29,18 @@
 if Version(jax.__version__) <= Version("0.10.0"):
   _OUT_KW = "out_shape"
   _SCRATCH_KW = "scratch_shapes"
+  _COMPILER_PARAMS = {
+      "use_tc_tiling_on_sc": True,
+      "disable_bounds_checks": True,
+  }
 else:
   _OUT_KW = "out_type"
   _SCRATCH_KW = "scratch_types"
+  _COMPILER_PARAMS = {
+      "use_tc_tiling_on_sc": True,
+      "disable_bounds_checks": True,
+      "needs_layout_passes": False,
+  }
 
 
 # ceil up to the nearest multiple of b.
@@ -488,9 +497,8 @@ def ragged_gather_reduce(
           num_row_partitions=num_rows_partitions,
           num_column_partitions=num_column_partitions,
       ),
-      compiler_params=pltpu.CompilerParams(
-          use_tc_tiling_on_sc=True,
-          disable_bounds_checks=True,
+      compiler_params=pltpu.CompilerParams(  # pytype: disable=wrong-keyword-args
+          **_COMPILER_PARAMS,
       ),
       mesh=vector_mesh,
       name="sc_ragged_gather_reduce",