Which part of the training data to read.

    If there are multiple parallel calls to input_fn (multiple TPU hosts),
    then we want each one to read from a separate partition of the training
    data.

    Args:
      mode: tf.estimator.ModeKeys
      config: RunConfig
      params: A dict that contains parameters.
    Returns:
      partition_id: an integer
      num_partitions: an integer

def _dataset_partition(self, mode, config, params):
    """Which part of the training data to read.

    If there are multiple parallel calls to input_fn (multiple TPU hosts),
    then we want each one to read from a separate partition of the training
    data.

    Args:
      mode: tf.estimator.ModeKeys
      config: RunConfig
      params: A dict that contains parameters.
    Returns:
      partition_id: an integer
      num_partitions: an integer
    """
    if mode != tf.estimator.ModeKeys.TRAIN or not hasattr(config, "tpu_config"):
      # Reset in the case when using TPU but alternating TRAIN and EVAL.
      self._next_partition_id = 0
      return 0, 1
    phift = config.tpu_config.per_host_input_for_training
    # This is the mesh-tensorflow case.
    if (hasattr(tpu_config.InputPipelineConfig, "BROADCAST") and
        phift == tpu_config.InputPipelineConfig.BROADCAST):
      return 0, 1
    if phift:
      num_hosts = (params["context"].num_hosts if "context" in params
                   else config.tpu_config.num_shards // 8)
      num_partitions = max(num_hosts, 1)
    else:
      num_partitions = config.tpu_config.num_shards
    partition_id = getattr(self, "_next_partition_id", 0)
    self._next_partition_id = partition_id + 1
    tf.logging.info("num_partitions = %d partition_id = %d" %
                    (num_partitions, partition_id))
    assert partition_id < num_partitions
    return partition_id, num_partitions