code: def pnum_to_group(mesh_shape, group_dims, pnum):
    
    coord = pnum_to_processor_coordinates(mesh_shape, pnum)
    remaining_shape = Shape(
        [d for i, d in enumerate(mesh_shape) if i not in group_dims])
    remaining_coord = [d for i, d in enumerate(coord) if i not in group_dims]
    return processor_coordinates_to_pnum(remaining_shape, remaining_coord)
docs: Group number for grouped allreduce.
  
    Args:
      mesh_shape: a Shape
      group_dims: a list of integers (the dimensions reduced over)
      pnum: an integer
  
    Returns:
      an integer