code: def _split_into_groups(n, max_group_size, mesh_dim_size):
    
    if n % mesh_dim_size != 0:
      raise ValueError(
          "n=%d is not a multiple of mesh_dim_size=%d" % (n, mesh_dim_size))
    num_groups = max(1, n // max_group_size)
    while (num_groups % mesh_dim_size != 0 or n % num_groups != 0):
      num_groups += 1
    group_size = n // num_groups
    tf.logging.info(
        "_split_into_groups(n=%d, max_group_size=%d, mesh_dim_size=%d)"
        " = (num_groups=%d group_size=%d)" %
        (n, max_group_size, mesh_dim_size, num_groups, group_size))
    return num_groups, group_size
docs: Helper function for figuring out how to split a dimensino into groups.
  
    We have a dimension with size n and we want to split it into
    two dimensions: n = num_groups * group_size
  
    group_size should be the largest possible value meeting the constraints:
      group_size <= max_group_size
      (num_groups = n/group_size) is a multiple of mesh_dim_size
  
    Args:
      n: an integer
      max_group_size: an integer
      mesh_dim_size: an integer
  
    Returns:
      num_groups: an integer
      group_size: an integer
  
    Raises:
      ValueError: if n is not a multiple of mesh_dim_size