code: def fix_channel_mask_conflict(
      graph_module: torch.fx.GraphModule, masks: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """
      Fix the mask conflict between channel and group.
      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.
      """
      prune_axis = detect_mask_prune_dim(graph_module, masks)
      prune_type = detect_channel_prune_type(graph_module, masks)
      print(prune_axis, prune_type)
  
      dependency_sets = build_channel_dependency(
          graph_module, prune_axis=prune_axis, prune_type=prune_type
      )
  
      sum_idx = (1, 2, 3) if prune_axis == 0 else (0, 2, 3)
      (_, _tmp_tensor) = list(masks.items())[0]
      device = list(_tmp_tensor.values())[0].device
  
      for d_set in dependency_sets:
          if len(d_set) <= 1:
              continue
          # channel_masks is a list, each element is None or a vector, for example:
          # [[0, 1, 1, 0, 0], [0, 0, 1, 1, 0], None], None means no channel
          # is pruned.
          channel_masks = []
          fine_grained = False
          for node in d_set:
              name = node.target
              if name in masks and "weight" in masks[name]:
                  sub_module = graph_module.get_submodule(name)
                  assert sub_module is not None
                  mask = masks[name]["weight"]
                  if isinstance(sub_module, torch.nn.Conv2d):
                      channel_mask = (
                          mask.abs().sum(sum_idx) != 0
                      ).int()
                      if prune_axis == 1:
                          channel_mask = channel_mask.repeat(
                              sub_module.groups
                          )
                      channel_masks.append(channel_mask)
                      if (
                          channel_mask.sum()
                          * (mask.numel() / mask.shape[prune_axis])
                      ).item() != (mask > 0).sum().item():
                          fine_grained = True
                  elif isinstance(sub_module, torch.nn.Linear):
                      if prune_axis == 1:
                          channel_masks.append(
                              (mask.abs().sum(0) != 0).int()
                          )
                      else:
                          channel_masks.append(
                              (mask.abs().sum(1) != 0).int()
                          )
                  elif isinstance(sub_module, torch.nn.Embedding):
                      if prune_axis == 0:
                          channel_masks.append(
                              (mask.abs().sum(0) != 0).int()
                          )
                  elif isinstance(sub_module, torch.nn.BatchNorm2d):
                      channel_masks.append(mask.int())
                  elif isinstance(sub_module, torch.nn.ConvTranspose2d):
                      # convtranspose have difference memory layout, so that we need create
                      # a tmp_sum_idx for conv_transpose
                      tmp_sum_idx = (
                          (0, 2, 3) if prune_axis == 0 else (1, 2, 3)
                      )
                      channel_mask = (
                          mask.abs().sum(tmp_sum_idx) != 0
                      ).int()
                      if prune_axis == 0:
                          channel_mask = channel_mask.repeat(
                              sub_module.groups
                          )
                      channel_masks.append(channel_mask)
                      if (
                          channel_mask.sum()
                          * (mask.numel() / mask.shape[1 - prune_axis])
                      ).item() != (mask > 0).sum().item():
                          fine_grained = True
                  else:
                      raise RuntimeError(
                          f"unsupported module type: {type(sub_module).__name__}"
                      )
              else:
                  # no mask means not pruned, equivlent to full masks
                  channel_masks.append(None)
          if fine_grained:
              _logger.info("Fine-grained mask detected")
          if all(x is None for x in channel_masks):
              continue
          num_channels_list = [
              len(x) for x in channel_masks if x is not None
          ]
          # number of channels in same set should be identical
          assert len(set(num_channels_list)) == 1
          num_channels = num_channels_list[0]
  
          for i, dim_mask in enumerate(channel_masks):
              if dim_mask is None:
                  channel_masks[i] = (
                      torch.ones(num_channels).int().to(device)
                  )
  
          # merge masks with 'or'
          merged_channel_mask = channel_masks[0].clone()
          for i in range(1, len(channel_masks)):
              merged_channel_mask = (
                  (merged_channel_mask + channel_masks[i]) != 0
              ).int()
  
          merged_index = torch.nonzero(
              merged_channel_mask, as_tuple=True
          )[0]
  
          for node in d_set:
              name = node.target
              if name not in masks or "weight" not in masks[name]:
                  assert all(merged_channel_mask)
                  continue
              orig_mask = masks[name]["weight"]
              sub_module = graph_module.get_submodule(name)
              new_mask = torch.zeros_like(orig_mask)
              if isinstance(sub_module, torch.nn.Conv2d):
                  if prune_axis == 0:
                      new_mask[merged_index, :, :, :] = 1.0
                  else:
                      new_mask[
                          :,
                          torch.nonzero(
                              merged_channel_mask[: new_mask.shape[1]],
                              as_tuple=True,
                          )[0],
                          :,
                          :,
                      ] = 1.0
              elif isinstance(sub_module, torch.nn.Linear):
                  if prune_axis == 0:
                      new_mask[merged_index, :] = 1.0
                  elif prune_axis == 1:
                      new_mask[:, merged_index] = 1.0
              elif isinstance(sub_module, torch.nn.Embedding):
                  if prune_axis == 0:
                      new_mask[:, merged_index] = 1.0
              elif isinstance(sub_module, torch.nn.BatchNorm2d):
                  new_mask = merged_channel_mask.type_as(orig_mask)
              else:
                  raise RuntimeError(
                      f"unsupported module type: {type(sub_module).__name__}"
                  )
              masks[name]["weight"] = new_mask
              if (
                  "bias" in masks[name]
                  and masks[name]["bias"] is not None
              ):
                  if prune_axis == 0:
                      masks[name]["bias"] = merged_channel_mask.type_as(
                          masks[name]["bias"]
                      )
      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 3 for calling a global function with example usage: fix_channel_mask_conflict(graph_module, masks) and returns: masks