def cross_shard_concat(input_tensor):
  """Concatenates all replicas of `input_tensor` across all shardss.

  Args:
    input_tensor: <int32|float32> [local_batch_size, dim1, dim2, ...]

  Returns:
    padded_tensor: <int32|float32>
        [local_batch_size * num_shards, dim1, dim2, ...]
  """
  padded_tensor = cross_shard_pad(input_tensor)
  concat_tensor = tf.tpu.cross_replica_sum(padded_tensor)
  return concat_tensor
Concatenates all replicas of `input_tensor` across all shardss.

  Args:
    input_tensor: <int32|float32> [local_batch_size, dim1, dim2, ...]

  Returns:
    padded_tensor: <int32|float32>
        [local_batch_size * num_shards, dim1, dim2, ...]
  
language/common/utils/tpu_utils.py