Repository: google/jax

Issue:
`vmap(grad(all_gather))` errors
Repro:
```python
import functools
import jax
import jax.numpy as jnp
from jax import lax
jax.config.update('jax_traceback_filtering', 'off')

@functools.partial(jax.vmap, in_axes=(0, None), axis_name='a')
def run(x, y):
  def f(x, y):
    return lax.all_gather(x, 'a')[y].sum()
  return jax.grad(f)(x, y)
run(jnp.ones(5), jnp.ones(10, jnp.int32))
```

Traceback:
There are a lot of frames (45) so I kept in the last few:
```
/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, *args, **params)
    263         args, used_axis_names(self, params) if self._dispatch_on_params else None)
    264     tracers = map(top_trace.full_raise, args)
--> 265     out = top_trace.process_primitive(self, tracers, params)
    266     return map(full_lower, out) if self.multiple_results else full_lower(out)
    267 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    138           _main_trace_for_axis_names(self.main, core.used_axis_names(primitive, params))):
    139       frame = core.axis_frame(self.axis_name)
--> 140       val_out, dim_out = collective_rules[primitive](frame, vals_in, dims_in, **params)
    141     elif all(bdim is not_mapped for bdim in dims_in):
    142       return primitive.bind(*vals_in, **params)

/usr/local/lib/python3.7/dist-packages/jax/_src/lax/parallel.py in _all_to_all_batched_collective(frame, vals_in, dims_in, axis_name, split_axis, concat_axis, axis_index_groups)
    901   if not major_axes and not minor_axes:
    902     if split_axis == concat_axis:
--> 903       axis = split_axis + (d <= split_axis)
    904       d_pre_split = d
    905       x = _splitaxis(axis, frame.size, x)

TypeError: '<=' not supported between instances of 'NoneType' and 'int'
```

@apaszke 


Discussion:
Some further investigation indicates that the `all_to_all_p` axis primitive batcher does not currently handle unbatched  (`dim=None`) inputs.

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
@@ -896,6 +896,12 @@ def _all_to_all_batched_collective(axis_size, frame_name, _, vals_in, dims_in,
     raise NotImplementedError("Please open a feature request!")
   x, = vals_in
   d, = dims_in
+  if d is batching.not_mapped:
+    # TODO(sharadmv,apaszke): Remove this broadcast that comes from
+    # all_gather_transpose and instead avoid using all_to_all in
+    # all_gather_transpose.
+    x = lax.broadcast(x, (axis_size, *x.shape))
+    d = 0
   if isinstance(axis_name, (list, tuple)):
     pos = axis_name.index(frame_name)
     major_axes, minor_axes = axis_name[:pos], axis_name[pos + 1:]
@@ -1087,6 +1093,9 @@ def _all_gather_transpose_rule(cts, x, *, all_gather_dimension, axis_name, axis_
       cts, axis_name=axis_name, split_axis=all_gather_dimension,
       concat_axis=concat_axis, axis_index_groups=axis_index_groups),
       axis=concat_axis),)
+  # TODO(sharadmv,apaszke): re-enable this when we can properly detect
+  # replication.
+  # return (lax.dynamic_index_in_dim(cts, idx, axis=all_gather_dimension, keepdims=False) * axis_size,)
 
 def _all_gather_batcher(vals_in, dims_in, *, all_gather_dimension, axis_name, axis_index_groups, axis_size, tiled):
   if tiled:


Test patch: