code: def replace_groupnorm(norm: nn.GroupNorm, masks):
      """
      Parameters
      ----------
      norm : torch.nn.GroupNorm
          The group norm module to be replace
      masks : Tuple of the input masks, output masks and weight masks
          Tuple of the masks, for example
          ([input_m1, input_m2], [output_m], {'weight':weight_m})
  
      Returns
      -------
      torch.nn.GroupNorm
          The new group norm module
      """
      in_masks, output_mask, _ = masks
      assert isinstance(norm, nn.GroupNorm)
      in_mask = in_masks[0]
  
      # N, C, H, W
      _, remained_in = convert_to_coarse_mask(in_mask, 1)
      _, remained_out = convert_to_coarse_mask(output_mask, 1)
  
      assert len(remained_in.size()) == 1
      if remained_in.size(0) != remained_out.size(0):
          raise ShapeMisMatchError()
  
      ori_channel_step = norm.num_channels // norm.num_groups
      for groupid in range(norm.num_groups):
          in_start = groupid * ori_channel_step
          in_end = in_start + ori_channel_step
  
          new_channel_step = (
              torch.logical_and(
                  in_start <= remained_in,
                  remained_in < in_end,
              )
              .sum()
              .item()
          )
  
          # this group fully pruned
          if new_channel_step == 0:
              continue
  
          break
  
      new_groups = 0
  
      # Validate
      for groupid in range(norm.num_groups):
          in_start = groupid * ori_channel_step
          in_end = in_start + ori_channel_step
          num_item = (
              torch.logical_and(
                  in_start <= remained_in,
                  remained_in < in_end,
              )
              .sum()
              .item()
          )
          if num_item == 0:
              continue
  
          # check if the number of remained channel of each group are the same
          if num_item != new_channel_step:
              raise UnBalancedGroupError()
  
          new_groups += 1
  
      new_num_channels = remained_in.size()[0]
      new_module = nn.GroupNorm(
          new_groups,
          new_num_channels,
          eps=norm.eps,
          affine=norm.affine,
      )
      if new_module.affine:
          new_module.weight.data = torch.index_select(
              norm.weight.data,
              0,
              remained_in,
          )
          new_module.bias.data = torch.index_select(
              norm.bias.data,
              0,
              remained_in,
          )
      return new_module
  
file_path: nni/nni/compression/speedup/replacement.py
imports: import logging
  import math
  import torch
  import torch.nn
  from .error_code import EmptyLayerError, ShapeMisMatchError, InputsNumberError, OutputTypeError, UnBalancedGroupError
label_desc: Use image node_id 8 for calling a global function with example usage: replace_groupnorm(norm, masks) and returns: new_module