Repository: jax-ml/jax
Pull request #24922 (2024-11-15 21:41:04)

[sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot
[sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot



Patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
index ee3c929b2..a1d326162 100644
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -2475,17 +2475,15 @@ wrap_with_shard_to_full_op = partial(_wrap_with_spmd_op, "SPMDShardToFullShape")
 
 
 def lower_sharding_under_shit(ctx, op, aval, sharding_proto=None):
-  if sharding_proto is None:
-    proto = aval.sharding._to_xla_hlo_sharding(aval.ndim).to_proto()
-  else:
-    proto = sharding_proto
-  # TODO(yashkatariya): Setting all axes as unspecified should work even when
-  # any axes is Collective because that's what happens in partial auto shmap.
-  # Do that after tests for it exists.
-  unspecified_dims = (set(range(aval.ndim))
-                      if aval.sharding.mesh.are_all_axes_collective else None)
-  return wrap_with_sharding_op(
-      ctx, op, aval, proto, unspecified_dims=unspecified_dims)
+  # Don't emit a wsc under full manual mode to avoid increasing HLO size.
+  if aval.sharding.mesh._are_all_axes_collective:
+    return op
+  proto = (aval.sharding._to_xla_hlo_sharding(aval.ndim).to_proto()
+           if sharding_proto is None else sharding_proto)
+  # TODO(yashkatariya): Enable this
+  # unspecified_dims = (set(range(aval.ndim))
+  #                     if aval.sharding.mesh._any_axis_collective else None)
+  return wrap_with_sharding_op(ctx, op, aval, proto)
 
 
 def set_sharding(op, sharding: xc.OpSharding | sharding_impls.SdyArraySharding):
diff --git a/jax/_src/mesh.py b/jax/_src/mesh.py
index 6c6017c4b..a2ab261fa 100644
--- a/jax/_src/mesh.py
+++ b/jax/_src/mesh.py
@@ -426,7 +426,7 @@ class AbstractMesh:
     return self.size == 0
 
   @functools.cached_property
-  def are_all_axes_collective(self) -> bool:
+  def _are_all_axes_collective(self) -> bool:
     if self.axis_types is None:
       return False
     return all(t == AxisTypes.Collective for t in self.axis_types.keys())


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 7196a6335..be1f9cfc2 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -5208,8 +5208,8 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     arr2 = jax.device_put(np_inp, NamedSharding(mesh, P('x', 'y')))
 
     def g(x, y):
-      self.assertTrue(x.sharding.mesh.are_all_axes_collective)
-      self.assertTrue(y.sharding.mesh.are_all_axes_collective)
+      self.assertTrue(x.sharding.mesh._are_all_axes_collective)
+      self.assertTrue(y.sharding.mesh._are_all_axes_collective)
       return x * y
 
     @jax.jit
@@ -5232,8 +5232,8 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     arr2 = jax.device_put(np_inp.T, NamedSharding(mesh, P('y', 'x')))
 
     def g(x, y):
-      self.assertTrue(x.sharding.mesh.are_all_axes_collective)
-      self.assertTrue(y.sharding.mesh.are_all_axes_collective)
+      self.assertTrue(x.sharding.mesh._are_all_axes_collective)
+      self.assertTrue(y.sharding.mesh._are_all_axes_collective)
       allgatherd_y = jax.lax.all_gather(y, axis_name='x', axis=1, tiled=True)
       z = x @ allgatherd_y
       return jax.lax.psum(z, axis_name='y')