code: def fix_group_mask_conflict(
      graph_module: torch.fx.GraphModule, masks: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """
      Fix the mask conflict between group and channel.
      This function will modify the masks in-place.
  
      Parameters
      ----------
      graph_module
          The graph module.
      masks
          The masks of all modules.
  
      Returns
      -------
      masks
          The fixed masks.
      """
      group_dependency = build_group_dependency(graph_module)
  
      for node, (max_group, min_group) in group_dependency.items():
          layername = node.target
          if layername not in masks or "weight" not in masks[layername]:
              continue
  
          w_mask = masks[layername]["weight"]
          shape = w_mask.shape
          count = np.prod(shape[1:])
  
          all_ones = (
              (w_mask.flatten(1).sum(-1) == count)
              .nonzero()
              .squeeze(1)
              .tolist()
          )
          all_zeros = (
              (w_mask.flatten(1).sum(-1) == 0)
              .nonzero()
              .squeeze(1)
              .tolist()
          )
          if len(all_ones) + len(all_zeros) < w_mask.size(0):
              # In fine-grained pruning, skip this layer
              _logger.info(
                  "Layers %s using fine-grained pruning", layername
              )
              continue
          assert shape[0] % max_group == 0
  
          # Find the number of masked filter for each group (mini_masked).
          # Because we have to keep the pruned filter can still
          # be divided into the same number of groups, so we only can
          # prune mini_masked filters for each group.
          step = shape[0] / max_group
          group_masked = []
          for i in range(max_group):
              _start = step * i
              _end = step * (i + 1)
              _tmp_list = list(
                  filter(lambda x: _start <= x and x < _end, all_zeros)
              )
              group_masked.append(_tmp_list)
          mini_masked = min([len(x) for x in group_masked])
          need_unmask = set()
          for gm in group_masked:
              for i in range(mini_masked, len(gm)):
                  # To keep the output channel number still being divisible to
                  # groups, we set the masks of following filters to be zero.
                  pos = gm[i]
                  need_unmask.add(pos)
          step = shape[0] / min_group
          for i in range(min_group):
              _start = step * i
              _end = step * (i + 1)
              _tmp_list = list(
                  filter(lambda x: _start <= x and x < _end, all_zeros)
              )
              if len(_tmp_list) == step:
                  # if the whole group is removed, then we don't have to unmask for
                  # the filters in this group
                  for pos in _tmp_list:
                      if pos in need_unmask:
                          need_unmask.remove(pos)
          for pos in need_unmask:
              masks[layername]["weight"][pos] = torch.ones(shape[1:])
              if hasattr(masks[layername], "bias"):
                  masks[layername]["bias"][pos] = 1
      return masks
  
file_path: nni/nni/compression/speedup/mask_conflict.py
imports: import logging
  from typing import Dict
  import numpy
  import torch
  import torch.fx
  from .dependency import build_channel_dependency, build_group_dependency, build_weight_sharing_dependency
label_desc: Use image node_id 1 for calling a global function with example usage: fix_group_mask_conflict(graph_module, masks) and returns: masks