Repository: jax-ml/jax
Issue category: Feature Request

[shmap] shard_map collectives don't support axis_index_groups
I am primarily interested in collective matrix multiplication algorithms where `all_to_all` and `all_gather` are useful for overlapping reduce scatter permutes and matrix multiplication.

Collectives like `psum` may be hard to support since the replication rule doesn't make sense, but collectives like `all_gather`, `all_to_all`, and `psum_scatter` should be easy since they are per-device and the output is unreplicated.

Example of code that doesn't work:

```python
from functools import partial

import jax
from jax import numpy as jnp
from jax import sharding
from jax.experimental.shard_map import shard_map

mesh = sharding.Mesh(jax.devices(), ('x',))
x = jax.device_put(
    jnp.arange(64).reshape(8, 8),
    sharding.NamedSharding(mesh, sharding.PartitionSpec(None, 'x')),
)
assert x.addressable_data(0).shape == (8, 1)


@partial(
    shard_map,
    mesh=mesh,
    in_specs=(sharding.PartitionSpec(None, 'x'),),
    out_specs=sharding.PartitionSpec(None, 'x'),
)
def fwd(x):
  return jax.lax.psum_scatter(
      x,
      'x',
      scatter_dimension=0,
      axis_index_groups=((0, 2, 4, 6), (1, 3, 5, 7)),
      tiled=True,
  )


y = fwd(x)

# Expected:
# [[ 12  16  76  80 140 144 204 208]
#  [ 44  48 108 112 172 176 236 240]]
```

See https://github.com/google/jax/blob/bbeddbdcffd7fcd8e52ea9b708c01fdfd4c52bb0/jax/experimental/shard_map.py#L926

cc @mattjj 


Functions to edit:
jax/experimental/shard_map.py:_standard_collective_rewrite

Patch:
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index a987d8d0faf4..40c9c216215e 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -923,7 +923,6 @@ def _standard_collective_check(prim, mesh, x_rep, *, axis_name, **params):
 
 def _standard_collective_rewrite(prim, mesh, in_rep, x, axis_name, **params):
   # The standard collective rewrite may insert a pbroadcast on the input.
-  if params.get('axis_index_groups') is not None: raise NotImplementedError
   axis_name = (axis_name,) if not isinstance(axis_name, tuple) else axis_name
   x_rep, = in_rep
   axis_name_set = set(axis_name)

Test patch: diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index d9776a3e6216..cec18fe47ae1 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -137,6 +137,28 @@ def fwd(a):
     for i, a_shard in enumerate(np.split(a, 2, axis=0)):
       self.assertAllClose(d.addressable_data(i), a_shard)
 
+  def test_all_gather_with_axis_index_groups(self):
+    mesh, a, _ = create_inputs(P('x', ('y', 'z')), P(None, None))
+
+    @jax.jit
+    @partial(
+        shard_map,
+        mesh=mesh,
+        in_specs=(P('x', ('y', 'z')),),
+        out_specs=P('x', ('y', 'z')),
+    )
+    def fwd(a):
+      return lax.all_gather(
+          a, ('y', 'z'), axis_index_groups=((0, 1), (2, 3)), axis=-1, tiled=True
+      )
+
+    c = fwd(a)
+    self.assertEqual(c.addressable_data(0).shape, (4, 4))
+    for i, row_block in enumerate(np.split(a, 2, axis=0)):
+      for j, block in enumerate(np.split(row_block, 2, axis=-1)):
+        self.assertAllClose(c.addressable_data(4 * i + 2 * j), block)
+        self.assertAllClose(c.addressable_data(4 * i + 2 * j + 1), block)
+
   def test_matmul_partial(self):
     raise unittest.SkipTest("invalid replication asserted by out_spec?")
 
@@ -175,6 +197,39 @@ def fwd(a, b):
     self.assertEqual(d.addressable_data(0).shape, (1, 8))
     self.assertAllClose(expected[:4] + expected[4:], d)
 
+  def test_reduce_scatter_with_axis_index_groups(self):
+    axis_index_groups = ((0, 2, 4, 6), (1, 3, 5, 7))
+    mesh, a, _ = create_inputs(P(None, ('x', 'y', 'z')), P(None, None))
+    assert a.addressable_data(0).shape == (8, 1)
+
+    @jax.jit
+    @partial(
+        shard_map,
+        mesh=mesh,
+        in_specs=(P(None, ('x', 'y', 'z')),),
+        out_specs=P(None, ('x', 'y', 'z')),
+    )
+    def fwd(a):
+      return lax.psum_scatter(
+          a,
+          ('x', 'y', 'z'),
+          scatter_dimension=0,
+          axis_index_groups=axis_index_groups,
+          tiled=True,
+      )
+
+    c = fwd(a)
+
+    self.assertEqual(c.addressable_data(0).shape, (2, 1))
+
+    sum_of_even_columns = np.sum(a[..., axis_index_groups[0]], -1)
+    for i, sums in enumerate(np.split(sum_of_even_columns, 4, 0)):
+      self.assertAllClose(np.squeeze(c.addressable_data(2 * i), -1), sums)
+
+    sum_of_odd_columns = np.sum(a[..., axis_index_groups[1]], -1)
+    for i, sums in enumerate(np.split(sum_of_odd_columns, 4, 0)):
+      self.assertAllClose(np.squeeze(c.addressable_data(2 * i + 1), -1), sums)
+
   def test_collective_permute(self):
     devices = np.array(jax.devices()[:8]) # Take up to 8 devices
     mesh = Mesh(devices, axis_names=('x'))
@@ -266,6 +321,44 @@ def fwd(a):
     c = fwd(a)
     assert (c == jnp.reshape(a.T, (1, 64))).all()
 
+  def test_all_to_all_with_axis_index_groups(self):
+    mesh_axes = dict(x=4)
+    devices = np.array(jax.devices()[: np.prod(tuple(mesh_axes.values()))])
+    mesh = Mesh(
+        devices.reshape(tuple(mesh_axes.values())),
+        axis_names=tuple(mesh_axes.keys()),
+    )
+    a = jax.device_put(
+        jnp.arange(4 * 4).reshape((4, 4)),
+        jax.sharding.NamedSharding(mesh, P('x', None)),
+    )
+    self.assertEqual(a.addressable_data(0).shape, (1, 4))
+
+    @jax.jit
+    @partial(
+        shard_map,
+        mesh=mesh,
+        in_specs=(P('x', None),),
+        out_specs=P(None, 'x'),
+    )
+    def fwd(a):
+      return lax.all_to_all(
+          a,
+          'x',
+          split_axis=1,
+          concat_axis=0,
+          axis_index_groups=((0, 1), (2, 3)),
+          tiled=True,
+      )
+
+    c = fwd(a)
+
+    # Each shard corresponds to a quadrant rather than a row.
+    self.assertEqual(c.addressable_data(0).shape, (2, 2))
+    for i, row_block in enumerate(np.split(a, 2, axis=0)):
+      for j, block in enumerate(np.split(row_block, 2, axis=-1)):
+        self.assertAllClose(block, c.addressable_data(2 * i + j))
+
   def test_eager_repr(self):
     mesh = Mesh(np.array(jax.devices()[:4]).reshape(2, 2), ('x', 'y'))
     s = None