def cross_shard_pad(input_tensor):
  """Cross shard pad.

  Assuming `input_tensor` is replicated over different TPU cores across the
  zeroth dimension, this creates a global tensor with unique chunks per replica.
  This function only fills in the local `input_tensor` and pads the non-local
  part of the tensor with zeros. Does not actually do any cross-shard
  communication.

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

  Returns:
    padded_tensor: <int32|float32>
        [local_batch_size * num_shards, dim1, dim2, ...]
  """
  num_shards = num_tpu_shards()

  # [num_shards]
  local_mask = tf.equal(tf.range(num_shards), shard_id())
  local_mask = tf.cast(local_mask, input_tensor.dtype)

  tensor_shape = tensor_utils.shape(input_tensor)
  local_batch_size = tensor_shape[0]
  global_batch_size = num_shards * local_batch_size

  # [num_shards, 1, 1, ...]
  for _ in tensor_shape:
    local_mask = tf.expand_dims(local_mask, -1)

  # [num_shards, local_batch_size, input_tensor_dim1, ...]
  padded_tensor = local_mask * tf.expand_dims(input_tensor, 0)

  # [global_batch_size, input_tensor_dim1, ...]
  padded_tensor = tf.reshape(padded_tensor,
                             [global_batch_size] + tensor_shape[1:])
  return padded_tensor
Cross shard pad.

  Assuming `input_tensor` is replicated over different TPU cores across the
  zeroth dimension, this creates a global tensor with unique chunks per replica.
  This function only fills in the local `input_tensor` and pads the non-local
  part of the tensor with zeros. Does not actually do any cross-shard
  communication.

  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