import logging


def _num_groups(num_tokens: int,
                max_group_size: int,
                num_experts: int,
                num_expert_replicas: int,
                strict_group_size: bool = False) -> int:
  """Returns the number of token routing groups.

  Note: For pjit-based training, all quantities are global.

  We select the smallest num_groups such that:
  - num_groups >= num_tokens / max_group_size (ensuring the group size is no
    larger than max_group_size),
  - num_tokens % num_groups = 0 (ensuring that the group size evenly divides
    into the num_tokens),
  - num_groups % (num_expert_replicas * num_experts) = 0 (ensuring that number
    of groups can be split across the total number of experts).

  Args:
    num_tokens: Number of tokens from input batch.
    max_group_size: Maximum size of each token routing group. Actual group size
      may end up being smaller.
    num_experts: Total number of unique experts.
    num_expert_replicas: Number of copies of each expert.
    strict_group_size: If True, fail if unable to set the token group size equal
      to max_group_size.

  Returns:
    Number of token routing groups.

  Raises:
    ValueError if we cannot find a group_size satisfying the above requirements.
  """
  # For pjit-based partitioning, we manipulated arrays globally. The number of
  # experts must evenly divide the number of (global) groups.
  min_num_groups = num_tokens // max_group_size
  min_num_groups = max(min_num_groups, num_expert_replicas * num_experts)

  def viable(n):
    """Returns true iff n is a viable number of groups."""
    return num_tokens % n == 0 and n % (num_expert_replicas * num_experts) == 0

  # Increase the number of groups (and decrease the group size) until we have
  # a viable number of groups.
  num_groups = min_num_groups
  while num_groups < num_tokens and not viable(num_groups):
    num_groups += 1

  if num_tokens % num_groups > 0:
    raise ValueError(
        'Group size and the number of experts must divide evenly into the '
        f'global number of tokens, but num_tokens={num_tokens}, while '
        f'num_groups={num_groups} for max_group_size={max_group_size} '
        f'and num_experts={num_experts}, each with {num_expert_replicas} '
        'replicas')

  group_size = num_tokens // num_groups
  logging.info(
      'Selected group_size=%d and num_groups=%d for input num_tokens=%d, '
      'max_group_size=%d, num_experts=%d and num_expert_replicas=%d',
      group_size, num_groups, num_tokens, max_group_size, num_experts,
      num_expert_replicas)

  if strict_group_size and group_size != max_group_size:
    raise ValueError(
        f'Selected group_size={group_size} is less than the '
        f'max_group_size={max_group_size}. Exiting because strict mode is '
        'active (strict_group_size=True)')

  return num_groups