def _get_batch_size(params, hparams, config):
  """Batch size determined by params dict, HParams, and RunConfig."""
  # If params specifies batch size, use that. TPUEstimator passes batch size in
  # params.
  batch_size = params and params.get("batch_size")

  # If not set, then we're running on CPU/GPU, so use the batch size from the
  # hparams, and multiply by the number of data shards.
  if not batch_size:
    batch_size = hparams.tpu_batch_size_per_shard
    if config:
      batch_size *= config.t2t_device_info["num_shards"]

  return batch_size