Why does check_rep fail?

[eitanporat]

from the JAX Scaling Book "Programming TPUs in JAX"
Problem 1

This is my solution

from jax.experimental.shard_map import shard_map as shmap

@functools.partial(shmap, check_rep=False, mesh=mesh, in_specs=shd.PartitionSpec('X', 'Y'), out_specs=shd.PartitionSpec(None, 'Y'))
def average_shard(x):
  assert x.shape == (S // X, D // Y)
  local_mean = x.mean(axis=0)
  replicated = jax.lax.all_gather(local_mean, axis_name='X')
  return replicated

average_A_shmap = average_shard(A)

It works with check_rep=False but without it it fails raising the error

ValueError: shard_map applied to the function 'average_shard' was given out_specs which require replication which can't be statically inferred given the mesh:

The mesh given has shape (2, 4) with corresponding axis names ('X', 'Y').

out_specs is PartitionSpec(None, 'Y') which implies that the corresponding output value is replicated across mesh axis 'X', but could not infer replication over any axes

Check if these output values are meant to be replicated over those mesh axes. If not, consider revising the corresponding out_specs entries. If so, consider disabling the check by passing the check_rep=False argument to shard_map.

What am I doing wrong? Doesn't AllGather replicate the results along the specified axis and thus the output have sharding (None, Y); as in the X axis is replicated and not sharded?

Running a small test with check_rep=False and looking at the shards they appear to be replicated

[Shard(device=TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), index=(slice(None, None, None), slice(0, 2, None)), replica_id=0, data=[[12. 13.]
  [44. 45.]]),
 Shard(device=TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), index=(slice(None, None, None), slice(2, 4, None)), replica_id=0, data=[[14. 15.]
  [46. 47.]]),
 Shard(device=TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), index=(slice(None, None, None), slice(4, 6, None)), replica_id=0, data=[[16. 17.]
  [48. 49.]]),
 Shard(device=TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), index=(slice(None, None, None), slice(6, 8, None)), replica_id=0, data=[[18. 19.]
  [50. 51.]]),
 Shard(device=TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), index=(slice(None, None, None), slice(0, 2, None)), replica_id=1, data=[[12. 13.]
  [44. 45.]]),
 Shard(device=TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1), index=(slice(None, None, None), slice(2, 4, None)), replica_id=1, data=[[14. 15.]
  [46. 47.]]),
 Shard(device=TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), index=(slice(None, None, None), slice(4, 6, None)), replica_id=1, data=[[16. 17.]
  [48. 49.]]),
 Shard(device=TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), index=(slice(None, None, None), slice(6, 8, None)), replica_id=1, data=[[18. 19.]
  [50. 51.]])]

[sssshhhhhh]

Had the same issue, tldr use all_gather_invariant (recently added and undocumented). all_gather is varying for efficiency as the transpose of psum_scatter explained here https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html

I needed an all_gather instead of just using out_specs because the jit decided to spam collective permutes for whatever reason.