def batch_size_per_device(batch_size: int, num_devices: int):
  """Return the batch size per device."""
  per_device_batch_size, ragged = divmod(batch_size, num_devices)

  if ragged:
    msg = 'batch size must be divisible by num devices, got {} and {}.'
    raise ValueError(msg.format(per_device_batch_size, num_devices))

  return per_device_batch_size