Repository: google/jax
Issue: pmap all_to_all shape error

On an 8-device machine, if I run a trivial pmap + all_to_all with concat_axis=1:
```
x = np.ones((8, 8, 64))
y = jax.pmap(lambda x: lax.all_to_all(x, axis_name='i', split_axis=0, concat_axis=1), axis_name='i')(x)
print(y.shape, onp.array(y).shape)  # --> returns (8, 8, 64) (8, 1, 512) !!!
```
whereas
```
x = np.ones((8, 8, 64))
y = jax.pmap(lambda x: lax.all_to_all(x, axis_name='i', split_axis=0, concat_axis=0), axis_name='i')(x)
print(y.shape, onp.array(y).shape)  # --> returns (8, 8, 64) (8, 8, 64)  no shape disagreement
```
even just trying to print the repr for the former case `y` yields the error `<repr(<jax.interpreters.pxla.ShardedDeviceArray at 0x7fe547474d08>) failed: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 0: want f32[8,8,64]{2,1,0}, got f32[8,1,512]{2,0,1}>`

Changed files:
jax/lax/lax_parallel.py, tests/pmap_test.py

diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
index 03488bd22..11d75b5f9 100644
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -557,13 +557,25 @@ def _all_to_all_translation_rule(c, x, *, split_axis, concat_axis, axis_name,
   replica_groups = _replica_groups(axis_env, axis_name, None)
   if len(replica_groups[0]) == 1:
     return x
-  else:
+  elif platform == 'tpu':
     split_count = len(replica_groups[0])
     if not all(split_count == len(g) for g in replica_groups):
       raise ValueError('Replica groups must be equally sized')
     replica_groups_protos = xc.make_replica_groups(replica_groups)
-    return xops.AllToAll(x, split_axis, concat_axis, split_count,
-                         replica_groups_protos)
+    if concat_axis == split_axis:
+      return xops.AllToAll(x, split_axis, concat_axis, split_count,
+                           replica_groups_protos)
+    else:
+      if concat_axis < split_axis:
+        split_axis += 1
+      elif split_axis < concat_axis:
+        concat_axis += 1
+      x = xla.lower_fun(partial(lax.expand_dims, dimensions=(concat_axis,)), multiple_results=False)(c, x)
+      x = xops.AllToAll(x, split_axis, concat_axis, split_count, replica_groups_protos)
+      x = xla.lower_fun(partial(lax.squeeze, dimensions=(split_axis,)), multiple_results=False)(c, x)
+      return x
+  else:
+    raise NotImplementedError("all_to_all and pswapaxes only supported on TPU")
 
 def _all_to_all_split_axis_rule(vals, which_mapped, split_axis, concat_axis,
                                 axis_name):
@@ -585,8 +597,15 @@ def _moveaxis(src, dst, x):
   perm.insert(dst, src)
   return lax.transpose(x, perm)
 
+def _all_to_all_abstract_eval(x, axis_name, split_axis, concat_axis):
+  input_aval = raise_to_shaped(x)
+  shape = list(input_aval.shape)
+  size = shape.pop(split_axis)
+  shape.insert(concat_axis, size)
+  return ShapedArray(tuple(shape), input_aval.dtype, weak_type=False)
+
 all_to_all_p = core.Primitive('all_to_all')
-all_to_all_p.def_abstract_eval(lambda x, **params: raise_to_shaped(x))
+all_to_all_p.def_abstract_eval(_all_to_all_abstract_eval)
 xla.parallel_translations[all_to_all_p] = _all_to_all_translation_rule
 ad.deflinear(all_to_all_p, _all_to_all_transpose_rule)
 pxla.multi_host_supported_collectives.add(all_to_all_p)
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index dd9313489..641e1760a 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -197,6 +197,55 @@ class PmapTest(jtu.JaxTestCase):
     ans = f(x)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @parameterized.named_parameters(
+      {"testcase_name": f"_split={split_axis}_concat={concat_axis}",
+      "split_axis": split_axis, "concat_axis": concat_axis}
+      for split_axis, concat_axis in it.product(range(2), range(2)))
+  def testAllToAll(self, split_axis, concat_axis):
+    if jtu.device_under_test() != "tpu":
+      raise SkipTest("all_to_all not implemented on non-TPU platforms")
+    pmap_in_axis = 0
+    shape = (xla_bridge.device_count(),) * 3
+    x = np.arange(np.prod(shape)).reshape(shape)
+
+    @partial(pmap, axis_name='i')
+    def f(x):
+      return lax.all_to_all(x, 'i', split_axis, concat_axis)
+    y = f(x)
+    if pmap_in_axis <= split_axis:
+      split_axis += 1
+    ref = jnp.moveaxis(x, (pmap_in_axis, split_axis),
+                          (concat_axis + 1, 0))
+    self.assertAllClose(y, ref)
+
+  @parameterized.named_parameters(
+      {"testcase_name": f"_split={split_axis}_concat={concat_axis}",
+       "split_axis": split_axis, "concat_axis": concat_axis}
+      for split_axis, concat_axis in it.product(range(2), range(2)))
+  def testAllToAllSplitAxis(self, split_axis, concat_axis):
+    if jtu.device_under_test() != "tpu":
+      raise SkipTest("all_to_all not implemented on non-TPU platforms")
+    if xla_bridge.device_count() < 4:
+      raise SkipTest("test requires at least four devices")
+    pmap_in_axis = 0
+    shape = (4, 4, 4)
+    x = np.arange(np.prod(shape)).reshape(shape)
+
+    @partial(pmap, axis_name='i')
+    @partial(pmap, axis_name='j')
+    def f(x):
+      return lax.all_to_all(x, ('i', 'j'), split_axis, concat_axis)
+
+    unroll_shape = (2, 2, *shape[1:])
+    x_unroll = x.reshape(unroll_shape)
+    y_unroll = f(x_unroll)
+    y = y_unroll.reshape(shape)
+
+    if pmap_in_axis <= split_axis:
+      split_axis += 1
+    ref = jnp.moveaxis(x, (pmap_in_axis, split_axis),
+                          (concat_axis + 1, 0))
+    self.assertAllClose(y, ref)
 
   def testNestedBasic(self):
     f = lambda x: lax.psum(lax.psum(x, 'i'), 'j')