Add a config option to remove size 1 mesh axis from ShapedArray.sharding.

PiperOrigin-RevId: 804088314

Author: yashk2810

jax/_src/config.py

@@ -233,6 +233,7 @@ def trace_context():
           eager_constant_folding.value,
           numpy_dtype_promotion.value,
           default_device.value, random_seed_offset.value,
+          remove_size_one_mesh_axis_from_type.value,
           threefry_partitionable.value,
           threefry_gpu_kernel_lowering.value,
           use_direct_linearize.value,
@@ -1150,6 +1151,13 @@ def _safer_randint_deprecation(new_val):
           'DO NOT RELY ON THIS FLAG.'),
     include_in_jit_key=True)
 
+remove_size_one_mesh_axis_from_type = bool_state(
+    name='jax_remove_size_one_mesh_axis_from_type',
+    default=False,
+    upgrade=True,
+    help="Removes mesh axes of size 1 from ShapedArray.sharding",
+    include_in_jit_key=True)
+
 # TODO make it so people don't use this, this is internal...
 _check_vma = bool_state(
     name='check_vma',

jax/_src/core.py

@@ -1684,7 +1684,7 @@ def mem_space_to_kind(mem_space: MemorySpace) -> str:
     assert False, "unreachable"
 
 
-@cache(max_size=4096, trace_context_in_key=False)
+@cache(max_size=4096, trace_context_in_key=True)
 def update_aval_with_sharding(aval, sharding):
   if isinstance(sharding, NamedSharding):
     return aval.update(
@@ -2082,6 +2082,17 @@ def modify_spec_for_auto_manual(spec, mesh) -> P:
                  if mesh._name_to_type[u] == AxisType.Explicit}
   return P(*new_spec, unreduced=new_unreduced, reduced=new_reduced)
 
+def remove_size_one_mesh_axis(spec, mesh) -> P:
+  new_spec = []  # type: ignore
+  for s in spec:
+    if s is None:
+      new_spec.append(s)  # type: ignore
+    elif isinstance(s, tuple):
+      new_spec.append(tuple(i for i in s if mesh.shape[i] != 1))
+    else:
+      new_spec.append(None if mesh.shape[s] == 1 else s)  # type: ignore
+  return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced)
+
 def _maybe_modify_sharding(sharding, ndim):
   if len(sharding.spec) == 0 or all(s is None for s in sharding.spec):
     out = sharding
@@ -2090,6 +2101,8 @@ def _maybe_modify_sharding(sharding, ndim):
   else:
     out = sharding.update(spec=modify_spec_for_auto_manual(
         sharding.spec, sharding.mesh))
+  if config.remove_size_one_mesh_axis_from_type.value:
+    out = out.update(spec=remove_size_one_mesh_axis(out.spec, out.mesh))
   if len(out.spec) != ndim:
     out = _make_lengths_same(out, ndim)
   return out
@@ -2108,7 +2121,7 @@ def _check_divisibility(sharding, shape):
           f" {size} times, but does not evenly divide the dimension size {sh}."
           f" Got shape: {shape} and sharding {sharding}")
 
-@cache(max_size=4096, trace_context_in_key=False)
+@cache(max_size=4096, trace_context_in_key=True)
 def get_sharding(sharding, shape):
   """Modifies and checks the sharding.
 

jax/_src/pjit.py

@@ -2725,6 +2725,8 @@ def check_shardings_are_auto(s: Sharding) -> None:
 def assert_shardings_equal(x_aval, user_sharding: NamedSharding):
   x_spec = x_aval.sharding.spec
   user_spec = user_sharding.spec._normalized_spec_for_aval(x_aval.ndim)
+  if config.remove_size_one_mesh_axis_from_type.value:
+    user_spec = core.remove_size_one_mesh_axis(user_spec, user_sharding.mesh)
   for x, s in zip(x_spec, user_spec):
     if s is PartitionSpec.UNCONSTRAINED:
       continue

tests/pjit_test.py

@@ -8658,6 +8658,29 @@ def f(x):
         f(inp)
     self.assertEqual(tracing_count(), 2)  # twice for f
 
+  @parameterized.named_parameters(
+      ('1', P('x', 'y'), P('x', None)),
+      ('2', P(('x', 'y')), P('x', None)),
+      ('3', P('y'), P(None, None)),
+  )
+  @config.remove_size_one_mesh_axis_from_type(True)
+  @jtu.with_explicit_mesh((2, 1), ('x', 'y'))
+  def test_remove_size_one_mesh_axis(self, arr_spec, type_spec, mesh):
+    arr = jax.device_put(np.arange(16).reshape(8, 2), arr_spec)
+
+    @jax.jit
+    def f(x):
+      self.assertEqual(x.aval.sharding.spec, type_spec)
+      out = x * 2
+      self.assertEqual(out.aval.sharding.spec, type_spec)
+      # wsc should act as an assert.
+      out = with_sharding_constraint(out, arr_spec)
+      return out
+
+    out = f(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, type_spec))
+    self.assertArraysEqual(out, arr * 2)
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):