Merge branch 'main' into jax_0.8.0_py3.12_v7x

andersensam · andersensam · commit a8b019a4f1d0 · 2025-10-21T19:22:38.000Z
diff --git a/axlearn/common/compiler_options.py b/axlearn/common/compiler_options.py
@@ -457,7 +457,6 @@ def infer_xla_performance_flags(
                 xla_tpu_enable_all_gather_offload_tracing="true",
                 xla_tpu_enable_reduce_scatter_offload_tracing="true",
                 xla_tpu_enable_all_reduce_offload_tracing="true",
-                xla_tpu_enable_reduce_scatter_legalizer="true",
             )
             # The available bandwidth of non-native mesh shapes is half of that compared to
             # the native mesh shape. Specify the latency modifier so that the latency hiding

Original file line number	Diff line number	Diff line change
`@@ -457,7 +457,6 @@ def infer_xla_performance_flags(`
`457`	`457`	`xla_tpu_enable_all_gather_offload_tracing="true",`
`458`	`458`	`xla_tpu_enable_reduce_scatter_offload_tracing="true",`
`459`	`459`	`xla_tpu_enable_all_reduce_offload_tracing="true",`
`460`		`- xla_tpu_enable_reduce_scatter_legalizer="true",`
`461`	`460`	`)`
`462`	`461`	`# The available bandwidth of non-native mesh shapes is half of that compared to`
`463`	`462`	`# the native mesh shape. Specify the latency modifier so that the latency hiding`