code: def _dataset_partition(self, mode, config, params):
      
      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
docs: 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