def unshard_in_pmap(tensors, n_shards):
  """Unshard tensors that were sharded into n_shards (call inside pmap)."""
  groups = [[n_shards * i + d for d in range(n_shards)]
            for i in range(fastmath.global_device_count() // n_shards)]
  def _unshard_fn(x):
    y = jax.lax.all_gather(x, 'batch', axis_index_groups=groups)
    split_y = jnp.split(y, n_shards, axis=0)
    split_y = [jnp.squeeze(sy, axis=0) for sy in split_y]
    axis = _axis_to_shard_heuristic(split_y[0].shape)
    return jnp.concatenate(split_y, axis=axis)
  try:
    jax.lax.axis_index('batch')  # will throw if not in pmap, e.g., on init
    res = fastmath.nested_map(_unshard_fn, tensors)
    return res, True
  except NameError:  # thrown from axis_index above
    return tensors, False
Unshard tensors that were sharded into n_shards (call inside pmap).
trax/layers/base.py