Repository: google/jax

Issue:
grad of shard_map of lax.all_to_all crashes
### Description

The forward transformation works, but not the backwards pass: 
```python
import functools
import chex
import jax
from jax.experimental import shard_map
from jax import lax
import jax.numpy as jnp
import numpy as np

chex.set_n_cpu_devices(4)

P = jax.sharding.PartitionSpec
shmap = shard_map.shard_map

mesh = jax.sharding.Mesh(jax.devices(), axis_names=['x'])

@functools.partial(
    shmap, mesh=mesh, in_specs=(P('x', None),), out_specs=P(None, 'x')
)
def reshard(x):
  return lax.all_to_all(x, 'x', split_axis=1, concat_axis=0, tiled=True)

def loss(x):
  return reshard(x).sum()

in_sharding = jax.sharding.NamedSharding(mesh, P('x', None))
x = jax.device_put(np.arange(64).reshape(8, 8), in_sharding)
print(x.sharding)  # NamedSharding(mesh={'x': 4}, spec=PartitionSpec('x', None))

y = reshard(x)
print(y.sharding)  # NamedSharding(mesh={'x': 4}, spec=PartitionSpec(None, 'x'))

np.testing.assert_array_equal(y, x)

print(jax.grad(loss)(1.0 * x))
```
`ValueError: all_to_all requires the size of the mapped axis axis_name to equal x.shape[split_axis], but they are 4 and 8 respectively.`

<details>

```
ValueError                                Traceback (most recent call last)
<ipython-input-72-29f919795aa1> in <module>()
     32 np.testing.assert_array_equal(y, x)
     33 
---> 34 print(jax.grad(loss)(1.0 * x))  # errors

third_party/py/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    164     __tracebackhide__ = True
    165     try:
--> 166       return fun(*args, **kwargs)
    167     except Exception as e:
    168       mode = _filtering_mode()

third_party/py/jax/_src/api.py in grad_f(*args, **kwargs)
    644   @api_boundary
    645   def grad_f(*args, **kwargs):
--> 646     _, g = value_and_grad_f(*args, **kwargs)
    647     return g
    648 

third_party/py/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    164     __tracebackhide__ = True
    165     try:
--> 166       return fun(*args, **kwargs)
    167     except Exception as e:
    168       mode = _filtering_mode()

third_party/py/jax/_src/api.py in value_and_grad_f(*args, **kwargs)
    726     _check_scalar(ans)
    727     tree_map(partial(_check_output_dtype_grad, holomorphic), ans)
--> 728     g = vjp_py(lax_internal._one(ans))
    729     g = g[0] if isinstance(argnums, int) else g
    730     if not has_aux:

third_party/py/jax/_src/tree_util.py in __call__(self, *args, **kw)
    301 
    302   def __call__(self, *args, **kw):
--> 303     return self.fun(*args, **kw)
    304 
    305   def __hash__(self):

third_party/py/jax/_src/api.py in _vjp_pullback_wrapper(name, cotangent_dtypes, cotangent_shapes, io_tree, fun, *py_args_)
   2102           "must be the same as the shape of corresponding primal input "
   2103           f"{ct_shape}.")
-> 2104   ans = fun(*args)
   2105   return tree_unflatten(out_tree, ans)
   2106 

third_party/py/jax/_src/tree_util.py in __call__(self, *args, **kw)
    301 
    302   def __call__(self, *args, **kw):
--> 303     return self.fun(*args, **kw)
    304 
    305   def __hash__(self):

third_party/py/jax/_src/interpreters/ad.py in unbound_vjp(pvals, jaxpr, consts, *cts)
    144     cts = tuple(ct for ct, pval in zip(cts, pvals) if not pval.is_known())
    145     dummy_args = [UndefinedPrimal(v.aval) for v in jaxpr.invars]
--> 146     arg_cts = backward_pass(jaxpr, reduce_axes, True, consts, dummy_args, cts)
    147     return map(instantiate_zeros, arg_cts)
    148 

third_party/py/jax/_src/interpreters/ad.py in backward_pass(jaxpr, reduce_axes, transform_stack, consts, primals_in, cotangents_in)
    251               reduce_axes, cts_in, *invals, **eqn.params)
    252         else:
--> 253           cts_out = get_primitive_transpose(eqn.primitive)(
    254               cts_in, *invals, **eqn.params)
    255         cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out

third_party/py/jax/experimental/shard_map.py in _shard_map_transpose(out_cts, jaxpr, mesh, in_names, out_names, check_rep, *args)
   1061     return tuple(names for names, nz in zip(in_names, nz_arg_cts()) if nz)
   1062 
-> 1063   out_flat = shard_map_p.bind(
   1064       fun_trans_flat, *all_args, mesh=mesh, in_names=tuple(new_in_names),
   1065       out_names_thunk=new_out_names_thunk, check_rep=check_rep)

third_party/py/jax/experimental/shard_map.py in bind(self, fun, mesh, in_names, out_names_thunk, check_rep, *args)
    345 
    346     tracers = map(top_trace.full_raise, args)
--> 347     outs = top_trace.process_shard_map(  # pytype: disable=attribute-error
    348         shard_map_p, fun, tracers, mesh=mesh, in_names=in_names,
    349         out_names_thunk=new_out_names_thunk, check_rep=check_rep)

third_party/py/jax/experimental/shard_map.py in _shard_map_impl(***failed resolving arguments***)
    552       t = main.with_cur_sublevel()
    553       in_tracers = map(partial(ShardMapTracer, t), in_rep, args)
--> 554       ans = fun.call_wrapped(*in_tracers)
    555       out_tracers = map(t.full_raise, ans)
    556       outs_, out_rep = unzip2((t.val, t.rep) for t in out_tracers)

third_party/py/jax/_src/linear_util.py in call_wrapped(self, *args, **kwargs)
    186 
    187     try:
--> 188       ans = self.f(*args, **dict(self.params, **kwargs))
    189     except:
    190       # Some transformations yield from inside context managers, so we have to

third_party/py/jax/experimental/shard_map.py in fun_trans(out_cts, args)
   1044         pe.close_jaxpr(jaxpr), map(ad.is_undefined_primal, args), False)
   1045     res_reshaped = core.jaxpr_as_fun(jaxpr_known)(*res)
-> 1046     out = ad.backward_pass(
   1047         jaxpr_unknown.jaxpr, (), False, (), (*res_reshaped, *undefs), out_cts
   1048     )

third_party/py/jax/_src/interpreters/ad.py in backward_pass(jaxpr, reduce_axes, transform_stack, consts, primals_in, cotangents_in)
    251               reduce_axes, cts_in, *invals, **eqn.params)
    252         else:
--> 253           cts_out = get_primitive_transpose(eqn.primitive)(
    254               cts_in, *invals, **eqn.params)
    255         cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out

third_party/py/jax/_src/interpreters/ad.py in linear_transpose2(transpose_rule, cotangent, *args, **kwargs)
    518 
    519 def linear_transpose2(transpose_rule, cotangent, *args, **kwargs):
--> 520   return Zero if type(cotangent) is Zero else transpose_rule(cotangent, *args, **kwargs)
    521 
    522 

third_party/py/jax/_src/lax/parallel.py in _all_to_all_transpose_rule(cts, x, axis_name, split_axis, concat_axis, axis_index_groups)
    993 
    994 def _all_to_all_transpose_rule(cts, x, axis_name, split_axis, concat_axis, axis_index_groups):
--> 995   return (all_to_all(
    996       cts,
    997       axis_name=axis_name,

third_party/py/jax/_src/lax/parallel.py in all_to_all(x, axis_name, split_axis, concat_axis, axis_index_groups, tiled)
    403     return result
    404 
--> 405   return tree_util.tree_map(bind, x)
    406 
    407 def axis_index(axis_name):

third_party/py/jax/_src/tree_util.py in tree_map(f, tree, is_leaf, *rest)
    208   leaves, treedef = tree_flatten(tree, is_leaf)
    209   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 210   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
    211 
    212 def build_tree(treedef: PyTreeDef, xs: Any) -> Any:

third_party/py/jax/_src/tree_util.py in <genexpr>(.0)
    208   leaves, treedef = tree_flatten(tree, is_leaf)
    209   all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
--> 210   return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
    211 
    212 def build_tree(treedef: PyTreeDef, xs: Any) -> Any:

third_party/py/jax/_src/lax/parallel.py in bind(x, split_axis, concat_axis)
    387         msg = ("all_to_all requires the size of the mapped axis axis_name to "
    388                "equal x.shape[split_axis], but they are {} and {} respectively.")
--> 389         raise ValueError(msg.format(group_size, x.shape[split_axis]))
    390       if split_axis < concat_axis:
    391         concat_axis += 1  # concat_axis gives a position _after_ split_axis is removed

ValueError: all_to_all requires the size of the mapped axis axis_name to equal x.shape[split_axis], but they are 4 and 8 respectively.
```
</details>

### What jax/jaxlib version are you using?

_No response_

### Which accelerator(s) are you using?

_No response_

### Additional system info

_No response_

### NVIDIA GPU info

_No response_


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
@@ -396,7 +396,8 @@ def bind(x, split_axis=split_axis, concat_axis=concat_axis):
         split_axis += 1   # we have a new axis before split_axis now
     result = all_to_all_p.bind(x, split_axis=split_axis, concat_axis=concat_axis,
                                axis_name=axis_name,
-                               axis_index_groups=axis_index_groups)
+                               axis_index_groups=axis_index_groups,
+                               tiled=tiled)
     if not tiled and split_axis != concat_axis:
       result = lax.squeeze(result, (split_axis,))
     return result
@@ -954,8 +955,10 @@ def _index_in_group(axis_name, axis_index_groups):
       slicing.dynamic_slice_in_dim(device_id_to_idx, cur_device_id, 1), [0])
 
 
-def _all_to_all_lowering(ctx, x, *,
-                         split_axis, concat_axis, axis_name, axis_index_groups):
+def _all_to_all_lowering(
+    ctx, x, *, split_axis, concat_axis, axis_name, axis_index_groups, tiled
+):
+  del tiled  # expand_dims and squeeze is done in `all_to_all` if `True`
   # Workaround for AllToAll not being implemented on CPU.
   replica_groups = _replica_groups(ctx.module_context.axis_env, axis_name,
                                    axis_index_groups)
@@ -985,15 +988,19 @@ def _all_to_all_lowering(ctx, x, *,
       replica_groups=_replica_groups_hlo(replica_groups),
       **other_args).results
 
-def _all_to_all_transpose_rule(cts, x, axis_name, split_axis, concat_axis, axis_index_groups):
+def _all_to_all_transpose_rule(
+    cts, x, axis_name, split_axis, concat_axis, axis_index_groups, tiled
+):
   return (all_to_all(
       cts,
       axis_name=axis_name,
       split_axis=concat_axis,
       concat_axis=split_axis,
-      axis_index_groups=axis_index_groups),)
+      axis_index_groups=axis_index_groups,
+      tiled=tiled),)
 
-def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis, axis_index_groups):
+def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis, axis_index_groups,
+                        tiled):
   x, = vals_in
   d, = dims_in
   result = all_to_all_p.bind(
@@ -1001,12 +1008,14 @@ def _all_to_all_batcher(vals_in, dims_in, *, axis_name, split_axis, concat_axis,
       axis_name=axis_name,
       split_axis=split_axis + (d <= split_axis),
       concat_axis=concat_axis + (d <= concat_axis),
-      axis_index_groups=axis_index_groups)
+      axis_index_groups=axis_index_groups,
+      tiled=tiled,
+  )
   return result, d
 
 def _all_to_all_batched_collective(axis_size, frame_name, _, vals_in, dims_in,
                                    axis_name, split_axis, concat_axis,
-                                   axis_index_groups):
+                                   axis_index_groups, tiled):
   if axis_index_groups is not None:
     raise NotImplementedError("Please open a feature request!")
   x, = vals_in
@@ -1041,7 +1050,8 @@ def _all_to_all_batched_collective(axis_size, frame_name, _, vals_in, dims_in,
   if major_axes:
     x = all_to_all_p.bind(x, axis_name=major_axes,
                           split_axis=split_axis, concat_axis=0,
-                          axis_index_groups=axis_index_groups)
+                          axis_index_groups=axis_index_groups,
+                          tiled=tiled)
   # Split out the local part into axis new_d (NOTE: d is already in axis 1)
   x = _splitaxis(split_axis, axis_size, x)
   new_d = split_axis
@@ -1050,7 +1060,8 @@ def _all_to_all_batched_collective(axis_size, frame_name, _, vals_in, dims_in,
   if minor_axes:
     x = all_to_all_p.bind(x, axis_name=minor_axes,
                           split_axis=split_axis, concat_axis=2,
-                          axis_index_groups=axis_index_groups)
+                          axis_index_groups=axis_index_groups,
+                          tiled=tiled)
 
   # Fold the chunk axes into a single one
   x = _foldaxis(0, _foldaxis(0, x))
@@ -1062,8 +1073,9 @@ def _all_to_all_batched_collective(axis_size, frame_name, _, vals_in, dims_in,
 
 
 def _all_to_all_effectful_abstract_eval(
-    x, axis_name, split_axis, concat_axis, axis_index_groups
+    x, axis_name, split_axis, concat_axis, axis_index_groups, tiled
 ):
+  del tiled  # expand_dims and squeeze is done in `all_to_all` if `True`
   if not isinstance(axis_name, (list, tuple)):
     axis_name = (axis_name,)
   input_aval = raise_to_shaped(x)


Test patch:
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -359,6 +359,40 @@ def fwd(a):
       for j, block in enumerate(np.split(row_block, 2, axis=-1)):
         self.assertAllClose(block, c.addressable_data(2 * i + j))
 
+  def test_all_to_all_grad(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(8 * 8, dtype=jnp.float32).reshape((8, 8)),
+        jax.sharding.NamedSharding(mesh, P('x', None)),
+    )
+    self.assertEqual(a.addressable_data(0).shape, (2, 8))
+
+    @jax.jit
+    @partial(
+        shard_map, mesh=mesh, in_specs=(P('x', None),), out_specs=P(None, 'x')
+    )
+    def fwd(x):
+      return lax.all_to_all(x, 'x', split_axis=1, concat_axis=0, tiled=True)
+
+    c = fwd(a)
+    self.assertEqual(c.addressable_data(0).shape, (8, 2))
+    self.assertAllClose(a, c)
+
+    @jax.jit
+    @partial(jax.grad, has_aux=True)
+    def loss_and_grad(x):
+      loss = fwd(x).sum() * 2
+      return loss, loss
+
+    grad, loss = loss_and_grad(a)
+    self.assertEqual(loss, 2 * sum(range(64)))
+    self.assertAllClose(grad, 2 * np.ones_like(a))
+
   def test_eager_repr(self):
     mesh = Mesh(np.array(jax.devices()[:4]).reshape(2, 2), ('x', 'y'))
     s = None