from typing import Sequence


def get_batch_dims(global_batch_size: int, device_count: int,
                   local_device_count: int) -> Sequence[int]:
  """Compute the batch dims for this host.

  The global_batch_size is the number of data samples that are optimized over
  in one step of the optimization. This value must be split up so that each
  individual device gets some share of the batch.

  When running with multiple devices, there may be multiple hosts, each
  with multiple local devices. Each host has a local copy of the program, and
  runs a local copy of the code. Each host must therefore use a batch size
  so that when all of the hosts run together, the total number of batched
  elements matches the global batch size. We do this by splitting up the global
  batch size evenly amongst all devices, and setting the batch size per host
  to the number of host devices times the device batch size.

  Args:
    global_batch_size: The target total batch size per optimization step.
    device_count: The total number of devices sharing computation per step.
    local_device_count: The number of devices available on the current host.

  Returns:
    The batch dimensions to use on the currently running host.
  """
  per_device_batch_size, remainder = divmod(global_batch_size, device_count)
  if remainder:
    raise ValueError(
        f'Cannot split batch of {global_batch_size} evenly across {local_device_count} devices.'
    )

  host_batch_dims = (local_device_count, per_device_batch_size)
  return host_batch_dims