code: def build_channel_dependency(
      graph_module: torch.fx.GraphModule,
      prune_type: str = "Filter",
      prune_axis: int = 0,
  ) -> List[Set[torch.fx.Node]]:
      """
      This model analyze the channel dependencies between the conv
      layers in a model.
  
      Parameters
      ----------
      graph_module : torch.fx.GraphModule
          The target graph and module.
      prune_type : str
          The channel pruning type. `Filter` prune the filter of the conv
          layer to prune the corresponding channels. `BatchNorm` prune the
          batchnorm layer to prune the corresponding channels.
      prune_axis : int
          The pruned axis of the conv layer. 1 for output channel, 0 for input channel.
  
      Returns
      -------
      dependency : List
          The channel dependency for the target graph.
      """
  
      if prune_type not in ["Filter", "BatchNorm"]:
          raise ValueError(
              "Unsupported prune type: {}".format(prune_type)
          )
  
      if prune_axis not in [0, 1]:
          raise ValueError(
              "Unsupported prune axis: {}. Support 0 for input channel and 1 for output channel.".format(
                  prune_axis
              )
          )
  
      graph = graph_module.graph
      target_types = [
          torch.nn.Conv2d,
          torch.nn.Linear,
          torch.nn.ConvTranspose2d,
          torch.nn.Embedding,
      ]
      if prune_type == "BatchNorm":
          target_types.append(torch.nn.BatchNorm2d)
  
      dependency = dict()
      for node in graph.nodes:
          d_set = set()
          node: torch.fx.Node
  
          # input channel dependency
          if prune_axis == 1:
              # Some pruners may prune the input channel of the convolutional
              # layers. While pruning the input channel of the convolutional layers,
              # the layers that share the same input tensor should prune the same
              # channels, and we say these layers that share the same input tensor/channel
              # has the input channel dependency. If we only prune the input channel of one
              # layer in the dependency set, there will be a shape conflict for the other
              # layers in the same dependency set, which may trigger a runtime error.
              # Here we judge whether the application will truncate the dependency by analyzing
              # whether the number of channels before and after the operation has changed.
              # If not, the input channel dependency will be passed to the following nodes.
  
              d_set = set(
                  find_adjacent_layers(
                      node, graph_module, target_types, "child"
                  )
              )
  
          # output channel dependency
          else:
              # Output channel dependency indicates the dependent layers when pruning the
              # output channles of conv layers (for example, L1FilterPruner).
  
              if node.op == "call_module":
                  submodule = graph_module.get_submodule(node.target)
                  # additional denpendency for (group number == output channel number) depth-wise conv:
                  if (
                      isinstance(
                          submodule,
                          (torch.nn.Conv2d, torch.nn.ConvTranspose2d),
                      )
                      and submodule.groups == submodule.out_channels
                  ) or (
                      isinstance(submodule, torch.nn.GroupNorm)
                      and submodule.num_groups == submodule.num_channels
                  ):
                      d_set = set(
                          [node]
                          + find_adjacent_layers(
                              node, graph_module, target_types, "parent"
                          )
                      )
  
              elif node.op == "call_function":
                  if node.target in CALL_FUNCTION_REDUCE:
                      # refer issue 4540 for more details. Multiplication actually
                      # will not introduce the channel dependency, cause the misaligned
                      # channels can propagate to each other. However, when one of the input
                      # tensor is from skip connection(residual), the channel propagation
                      # may be failed(the input is also used by another layer and cannot be
                      # pruned), in this case, we need to fix the conflict maunally.
                      d_set = set(
                          find_adjacent_layers(
                              node, graph_module, target_types, "parent"
                          )
                      )
  
                  elif node.target in CALL_FUNCTION_CONCAT:
                      # To determine if this cat operation will introduce channel
                      # dependency, we need the specific input parameters of the cat
                      # operation.
                      cat_dim = (
                          node.kwargs.get("dim", None)
                          if node.kwargs.get("dim", None) is not None
                          else node.args[1]
                      )
                      if cat_dim != 1:
                          d_set = set(
                              find_adjacent_layers(
                                  node,
                                  graph_module,
                                  target_types,
                                  "parent",
                              )
                          )
  
              elif node.op == "call_method":
                  if node.target in CALL_METHOD_REDUCE:
                      d_set = set(
                          find_adjacent_layers(
                              node, graph_module, target_types, "parent"
                          )
                      )
  
          # merge dependencies
          for parent in tuple(d_set):
              if parent in dependency:
                  d_set |= dependency[parent]
  
          for parent in d_set:
              dependency[parent] = d_set
  
      return convert_dependency_to_set(dependency)
  
file_path: nni/nni/compression/speedup/dependency.py
imports: import logging
  import operator
  from copy import deepcopy
  from typing import Any, Dict, List, Set
  from uuid import uuid4
  import numpy
  import torch
  import torch.fx
  from torch.fx.passes.shape_prop import TensorMetadata
  from nni.compression.base.config import trans_legacy_config_list, select_modules_by_config
label_desc: Use image node_id 8 for calling a global function with example usage: build_channel_dependency(graph_module, prune_type, prune_axis) and returns: convert_dependency_to_set