[SC] Specify tiling in mosaic kernel so it does not rely on global flag xla_tpu_use_tc_device_shape_on_sc.

megrez-yliu · Google-ML-Automation · commit 57a6b307f5d6 · 2026-06-02T10:18:41.000-07:00
PiperOrigin-RevId: 925433728
diff --git a/src/maxtext/kernels/gather_reduce_sc.py b/src/maxtext/kernels/gather_reduce_sc.py
@@ -1617,6 +1617,7 @@ def kernel_main(
     return mosaic.as_tpu_kernel(
         m,
         out_type=out_type,
+        tiling=mosaic.Tiling.COMPACT,
     )(
         *(
             [

Original file line number	Diff line number	Diff line change
`@@ -1617,6 +1617,7 @@ def kernel_main(`
`1617`	`1617`	`return mosaic.as_tpu_kernel(`
`1618`	`1618`	`m,`
`1619`	`1619`	`out_type=out_type,`
	`1620`	`+ tiling=mosaic.Tiling.COMPACT,`
`1620`	`1621`	`)(`
`1621`	`1622`	`*(`
`1622`	`1623`	`[`