def shard_id():
  """Get an integer scalar Tensor indicating the index of the current shard."""
  # Prevent the TPU compiler from rewriting this part of the graph.
  with tf.control_dependencies(None):
    return tpu_ops.tpu_replicated_input(list(range(num_tpu_shards())))
Get an integer scalar Tensor indicating the index of the current shard.
language/common/utils/tpu_utils.py