Repository: google/jax

Issue:
No batching rule for `all_to_all`
```py
import jax
from jax import lax, vmap, pmap
import jax.numpy as jnp

def f(x):
  tiled_x = jax.tree_map(lambda x: jnp.tile(x[None], [jax.device_count()] + [1] * len(x.shape)), x)
  all_x = lax.pswapaxes(tiled_x, 'batch', 0)
  return all_x
pmap(vmap(f), axis_name='batch')(jnp.arange(jax.device_count() * 5 * 2).reshape((jax.device_count(), 5, 2)))
```

Currently this throws an assertion error elsewhere but with #4140, it will correctly throw `NotImplementedError: Batching rule for 'all_to_all' not implemented`.


Discussion:

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -857,7 +857,7 @@ def _tuple_output(*args, **kwargs):
   ans = yield args, kwargs
   yield (ans,)
 
-def lower_fun(fun, multiple_results):
+def lower_fun(fun, multiple_results, parallel=False):
   # This function can only be used to lower functions that take JAX array types
   # as arguments (and e.g. don't accept unit values), because it assumes it can
   # map from XLA types to JAX types. In general that mapping is not possible (as
@@ -868,10 +868,14 @@ def lower_fun(fun, multiple_results):
   def f(c, *xla_args, **params):
     # TODO(mattjj): revise this 'calling convention'
     avals = [_array_aval_from_xla_shape(c.get_shape(x)) for x in xla_args]
+    if parallel:
+      axis_env = params.pop('axis_env')
+      del params['platform']
+    else:
+      axis_env = AxisEnv(1, (), (), None)
     wrapped_fun = lu.wrap_init(fun, params)
     if not multiple_results:
       wrapped_fun = _tuple_output(wrapped_fun)
-    axis_env = AxisEnv(1, (), (), None)
     if config.omnistaging_enabled:
       jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, avals)
       outs = jaxpr_subcomp(c, jaxpr, None, axis_env, _xla_consts(c, consts), '',
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -16,6 +16,7 @@
 """
 
 import collections
+import warnings
 
 import numpy as np
 
@@ -243,9 +244,7 @@ def pswapaxes(x, axis_name, axis):
       ``axis_name``.
 
   Returns:
-    Array(s) with shape ``np.insert(np.delete(x.shape, axis), axis, axis_size)``
-    where ``axis_size`` is the size of the mapped axis named ``axis_name`` in
-    the input ``x``.
+    Array(s) with the same shape as ``x``.
   """
   return all_to_all(x, axis_name, axis, axis)
 
@@ -551,13 +550,32 @@ def _ppermute_batcher(vals_in, dims_in, axis_size, axis_name, perm):
 batching.collective_rules[ppermute_p] = _ppermute_batcher
 
 
+def _moveaxis(src, dst, x):
+  perm = [i for i in range(x.ndim) if i != src]
+  perm.insert(dst, src)
+  return lax.transpose(x, perm)
+
+def _all_to_all_via_all_gather(x, *, axis_name, split_axis, concat_axis):
+  global_full = all_gather(x, axis_name)
+  idx = axis_index(axis_name)
+  local_slice = lax.dynamic_index_in_dim(global_full, idx, split_axis + 1, keepdims=False)
+  return _moveaxis(0, concat_axis, local_slice)
+
 def _all_to_all_translation_rule(c, x, *, split_axis, concat_axis, axis_name,
                                  axis_env, platform):
   # Workaround for AllToAll not being implemented on CPU.
   replica_groups = _replica_groups(axis_env, axis_name, None)
   if len(replica_groups[0]) == 1:
     return x
-  elif platform == 'tpu':
+  elif platform != 'tpu':
+    warnings.warn("all_to_all (and pswapaxes) are only implemented properly for TPUs. All other "
+                  "backends emulate it using a very slow and memory intensive algorithm, so expect "
+                  "significant slowdowns.")
+    lowering = xla.lower_fun(_all_to_all_via_all_gather, multiple_results=False, parallel=True)
+    return lowering(c, x,
+                    split_axis=split_axis, concat_axis=concat_axis, axis_name=axis_name,
+                    axis_env=axis_env, platform=platform)
+  else:
     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')
@@ -574,28 +592,25 @@ def _all_to_all_translation_rule(c, x, *, split_axis, concat_axis, axis_name,
       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):
-  assert tuple(which_mapped) == (True,)
-  x, = vals
-  # perform the communication to swap the hardware-mapped axes
-  stacked = all_to_all_p.bind(x, split_axis=split_axis + 1, concat_axis=0,
-                              axis_name=axis_name)
-  # transpose the newly mapped axis to the front, newly unmapped to concat_axis
-  out = _moveaxis(split_axis + 1, 0, stacked)
-  out = _moveaxis(1, concat_axis + 1, out)
-  return out, True
 
 def _all_to_all_transpose_rule(cts, axis_name, split_axis, concat_axis):
   return (all_to_all(cts, axis_name=axis_name, split_axis=concat_axis, concat_axis=split_axis),)
 
-def _moveaxis(src, dst, x):
-  perm = [i for i in range(x.ndim) if i != src]
-  perm.insert(dst, src)
-  return lax.transpose(x, perm)
+def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis):
+  x, = vals_in
+  d, = dims_in
+  if d <= split_axis:
+    split_axis += 1
+  if d <= concat_axis:
+    concat_axis += 1
+  # Note: At this point split_axis and concat_axis are adjusted to the extra
+  #       dimension and we have d != split_axis and d != concat_axis.
+  if split_axis < d < concat_axis:
+    d -= 1
+  elif concat_axis < d < split_axis:
+    d += 1
+  result = all_to_all_p.bind(x, axis_name=axis_name, split_axis=split_axis, concat_axis=concat_axis)
+  return result, d
 
 def _all_to_all_abstract_eval(x, axis_name, split_axis, concat_axis):
   input_aval = raise_to_shaped(x)
@@ -609,6 +624,7 @@ def _all_to_all_abstract_eval(x, axis_name, split_axis, concat_axis):
 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)
+batching.primitive_batchers[all_to_all_p] = _all_to_all_batcher
 
 
 def _expand(dim, size, index, x):


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -97,6 +97,8 @@ def tearDownModule():
 ignore_jit_of_pmap_warning = partial(
   jtu.ignore_warning, message=".*jit-of-pmap.*")
 
+ignore_slow_all_to_all_warning = partial(
+  jtu.ignore_warning, message="all_to_all.*expect significant slowdowns.*")
 
 class PmapTest(jtu.JaxTestCase):
   def _getMeshShape(self, device_mesh_shape):
@@ -143,6 +145,7 @@ def testGather(self):
     ans = f(x)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  @ignore_slow_all_to_all_warning()
   def testTrees(self):
     ptranspose = lambda x, axis_name: lax.all_to_all(x, axis_name, 0, 0)
     def protate(x, axis_name):
@@ -166,9 +169,9 @@ def protate(x, axis_name):
     assert_allclose(jax_f(lax.pmin)(x), np_f(np.min)(x))
     assert_allclose(jax_f(lax.psum)(x), np_f(np.sum)(x))
     assert_allclose(jax_f(lax.pmean)(x), np_f(np.mean)(x))
+    assert_allclose(jax_f(ptranspose)(x), np_transpose(x))
+    # NOTE: ppermute only supported on TPU.
     if jtu.device_under_test() not in ("cpu", "gpu"):
-      # NOTE: all-to-all and ppermute only supported on TPU.
-      assert_allclose(jax_f(ptranspose)(x), np_transpose(x))
       assert_allclose(jax_f(protate)(x), np_rotate(x))
 
   def testCollectivesWithTreesOfDifferentDtypes(self):
@@ -1630,6 +1633,77 @@ def f(x, y):
     x = jnp.ones((2, 2, 64, 64))
     self.assertAllClose(f(jax.pmap)(x, x), f(jax.vmap)(x, x))
 
+  @parameterized.named_parameters(
+      {"testcase_name": f"_split={split_axis}_concat={concat_axis}_vmap={vmap_axis}",
+       "split_axis": split_axis, "concat_axis": concat_axis, "vmap_axis": vmap_axis}
+      for split_axis, concat_axis, vmap_axis in it.product(range(3), range(3), range(4)))
+  @skipIf(not jax.config.omnistaging_enabled,
+          "vmap collectives only supported when omnistaging is enabled")
+  @ignore_slow_all_to_all_warning()
+  def testAllToAllInVmap(self, split_axis, concat_axis, vmap_axis):
+    def f(x):
+      return lax.all_to_all(x, 'i', split_axis=split_axis, concat_axis=concat_axis)
+
+    def adj(axis, hidden_axes):
+      for hax in sorted(hidden_axes):
+        if hax <= axis:
+          axis += 1
+      return axis
+
+    def reference(x, split_axis, concat_axis, vmap_axis):
+      pmap_axis = 0
+      vmap_axis = adj(vmap_axis, [pmap_axis])
+      ref = x
+
+      # Step 1.
+      # Adjust the split axis to the real tensor layout and move it to
+      # position 1. Since pmap_axis is always 0 we don't have to adjust it,
+      # but we do have to adjust vmap_axis.
+      split_axis = adj(split_axis, [pmap_axis, vmap_axis])
+      ref = jnp.moveaxis(ref, split_axis, pmap_axis + 1)
+      vmap_axis = vmap_axis + (0 if split_axis < vmap_axis else 1)
+      split_axis = pmap_axis + 1  # split_axes == 1
+
+      # Step 2.
+      # Now, we move pmap_axis to the position indicated by concat_axis.
+      concat_axis = adj(concat_axis, [pmap_axis, split_axis, vmap_axis]) - 1
+      ref = jnp.moveaxis(ref, pmap_axis, concat_axis)
+      pmap_axis = 0
+      vmap_axis = vmap_axis - (1 if concat_axis >= vmap_axis else 0)
+      del split_axis, concat_axis
+
+      # Step 3. vmap_axis always ends in position 1, since out_axes=0.
+      ref = jnp.moveaxis(ref, vmap_axis, 1)
+      return ref
+
+    def verify_ref():
+      # Both the reference and the real implementation of all_to_all batching involve
+      # some pretty complicated axis arithmetic, so it would be good to verify that it's
+      # not the case that the test passes because they're both incorrect. Fortunately, it
+      # is quite easy to write out the shape function for this code, and we know
+      # that it should be equivalent to a bunch of transposes, so the code below verifies
+      # that the reference puts the right dimensions in the right places. Note that we
+      # can't do the same comparison on f, since all_to_all wouldn't allow us to swap axes of
+      # different sizes.
+      start_shape = [2, 3, 4, 5, 6]
+      instance_shape = start_shape.copy()
+      pmap_dim_id = instance_shape.pop(0)
+      vmap_dim_id = instance_shape.pop(vmap_axis)
+      split_axis_id = instance_shape.pop(split_axis)
+      instance_shape.insert(concat_axis, pmap_dim_id)
+      expected_shape = (split_axis_id, vmap_dim_id, *instance_shape)
+
+      x = np.empty(start_shape)
+      self.assertEqual(reference(x, split_axis, concat_axis, vmap_axis).shape,
+                       expected_shape)
+
+    verify_ref()
+
+    shape = (jax.device_count(),) * 5
+    x = jnp.arange(np.prod(shape)).reshape(shape)
+    self.assertAllClose(pmap(vmap(f, in_axes=vmap_axis), axis_name='i')(x),
+                        reference(x, split_axis, concat_axis, vmap_axis))
+
 
 class PmapWithDevicesTest(jtu.JaxTestCase):