Repository: google/jax

Issue:
Implement tiled=True for all_gather inside a vmap  (to address NotImplementedError exception)
My scenario:
   I use an all_gather with tiled=True inside a pmap.  I hit an error which is nearly impossible to understand because pmap JITs its code.  So I tried the recommended strategy of substituting "vmap" for "pmap".  I then hit the following exception:     raise NotImplementedError("Please open a feature request!")

This exception is raised because "tiled" is set to True for my all_gather.  The exception is raised by the second line of _all_gather_batched_collective() in parallel.py





Discussion:

Resolving patch:
diff --git a/jax/_src/lax/parallel.py b/jax/_src/lax/parallel.py
--- a/jax/_src/lax/parallel.py
+++ b/jax/_src/lax/parallel.py
@@ -1089,12 +1089,10 @@ def _all_gather_transpose_rule(cts, x, *, all_gather_dimension, axis_name, axis_
       axis=concat_axis),)
 
 def _all_gather_batcher(vals_in, dims_in, *, all_gather_dimension, axis_name, axis_index_groups, axis_size, tiled):
-  if tiled:
-    raise NotImplementedError("Please open a feature request!")
   (x,), (d,) = vals_in, dims_in
   if d <= all_gather_dimension:
     all_gather_dimension += 1
-  else:
+  elif not tiled:  # Tiled all-gather doesn't modify the set of dimensions
     d += 1
   result = all_gather_p.bind(
       x,
@@ -1107,8 +1105,6 @@ def _all_gather_batcher(vals_in, dims_in, *, all_gather_dimension, axis_name, ax
 
 def _all_gather_batched_collective(frame_size, frame_name, _, vals_in, dims_in, all_gather_dimension, axis_name,
                                    axis_index_groups, axis_size, tiled):
-  if tiled:
-    raise NotImplementedError("Please open a feature request!")
   assert axis_index_groups is None, "axis_index_groups not supported in vmap"
   assert axis_size == frame_size, "axis size doesn't match"
   if not isinstance(axis_name, tuple):
@@ -1121,8 +1117,12 @@ def _all_gather_batched_collective(frame_size, frame_name, _, vals_in, dims_in,
     out_shape = list(np.shape(x))
     out_shape.insert(all_gather_dimension, axis_size)
     broadcast_dims = [i for i in range(len(out_shape)) if i != all_gather_dimension]
-    return lax.broadcast_in_dim(x, out_shape, broadcast_dims), batching.not_mapped
-  return _moveaxis(d, all_gather_dimension, x), batching.not_mapped
+    y = lax.broadcast_in_dim(x, out_shape, broadcast_dims)
+  else:
+    y = _moveaxis(d, all_gather_dimension, x)
+  if tiled:
+    y = _foldaxis(all_gather_dimension, y)
+  return y, batching.not_mapped
 
 all_gather_p = core.AxisPrimitive('all_gather')
 all_gather_p.def_abstract_eval(_all_gather_abstract_eval)


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -1193,6 +1193,22 @@ def f(x):
     res = vmap(vmap(f, axis_name='j'), axis_name='i', out_axes=None)(x)
     self.assertAllClose(res, x.T)
 
+  def testAllGatherTiled(self):
+    def f(x):
+      return lax.all_gather(x, axis_name='i', tiled=True)
+
+    x = jnp.arange(60).reshape((4, 3, 5))
+    res = vmap(f, axis_name='i', in_axes=(1,), out_axes=None)(x)
+    self.assertAllClose(res, x.transpose((1, 0, 2)).reshape(-1, 5))
+
+  def testBatchedAllGatherTiled(self):
+    def f(x):
+      return lax.all_gather(x, axis_name='i', tiled=True)
+
+    x = jnp.arange(60).reshape((4, 3, 5))
+    res = vmap(vmap(f, in_axes=1, out_axes=1), axis_name='i', in_axes=1, out_axes=None)(x)
+    self.assertAllClose(res, x.transpose((1, 0, 2)).reshape(-1, 5))
+
   def testAllGatherVjp(self):
     def f(x):
       return lax.all_gather(x, axis_name='i')