import logging


def get_batch_size(params, batch_size):
  """Get the correct batch size to use in an input_fn.

  TPUEstimator may provide different batch sizes if using it in
  contrib_tpu.InputPipelineConfig.PER_HOST_V2 setting.
  See comments about `per_host_input_for_training` in
  //third_party/tensorflow/python/tpu/tpu_config.py.

  Args:
    params: Hyperparameters passed to the input_fn. For TPUEstimator this
      includes 'batch_size'
    batch_size: Batch size provided to the input_fn.

  Returns:
    The batch size to invoke for the input_fn.
  """
  params_batch_size = params.get('batch_size') if params else None
  if params_batch_size and params_batch_size != batch_size:
    logging.info(
        'The input_fn has a batch_size set through `params`, as '
        'well as in the input generator. These batch sizes do '
        'not match. Using the batch size %d from params', params_batch_size)
    return params_batch_size
  return batch_size