import logging


def get_batch_size(num_replicas, num_samples, steps):
  """Calculate and return batch size for numpy inputs.

  Args:
    num_replicas: Number of devices over which the model input is distributed.
    num_samples: Total number of input samples in the input numpy arrays.
    steps: Number of steps that we run the model for.

  Returns:
    batch size used to create the Dataset object from the input numpy arrays.

  """
  if num_samples % steps != 0:
    logging.warning('The number of input samples %d is not evenly '
                    'divisible by the number of steps %d. '
                    'Some samples will not be processed as expected.' %
                    (num_samples, steps))
  global_batch_size = num_samples // steps
  if global_batch_size % num_replicas != 0:
    logging.warning('The total number of batches per step %d is not evenly '
                    'divisible by the number of replicas %d used in '
                    'DistributionStrategy. Some samples will not be processed '
                    'as expected.' %
                    (global_batch_size, num_replicas))
  return global_batch_size // num_replicas