code: def compute_vertex_channels(input_channels, output_channels, matrix):
      """
      This is (almost) copied from the original NAS-Bench-101 implementation.
  
      Computes the number of channels at every vertex.
  
      Given the input channels and output channels, this calculates the number of channels at each interior vertex.
      Interior vertices have the same number of channels as the max of the channels of the vertices it feeds into.
      The output channels are divided amongst the vertices that are directly connected to it.
      When the division is not even, some vertices may receive an extra channel to compensate.
  
      Parameters
      ----------
      in_channels : int
          input channels count.
      output_channels : int
          output channel count.
      matrix : np.ndarray
          adjacency matrix for the module (pruned by model_spec).
  
      Returns
      -------
      list of int
          list of channel counts, in order of the vertices.
      """
  
      num_vertices = np.shape(matrix)[0]
  
      vertex_channels = [0] * num_vertices
      vertex_channels[0] = input_channels
      vertex_channels[num_vertices - 1] = output_channels
  
      if num_vertices == 2:
          # Edge case where module only has input and output vertices
          return vertex_channels
  
      # Compute the in-degree ignoring input, axis 0 is the src vertex and axis 1 is
      # the dst vertex. Summing over 0 gives the in-degree count of each vertex.
      in_degree = np.sum(matrix[1:], axis=0)
      interior_channels = output_channels // in_degree[num_vertices - 1]
      correction = (
          output_channels % in_degree[num_vertices - 1]
      )  # Remainder to add
  
      # Set channels of vertices that flow directly to output
      for v in range(1, num_vertices - 1):
          if matrix[v, num_vertices - 1]:
              vertex_channels[v] = interior_channels
              if correction:
                  vertex_channels[v] += 1
                  correction -= 1
  
      # Set channels for all other vertices to the max of the out edges, going backwards.
      # (num_vertices - 2) index skipped because it only connects to output.
      for v in range(num_vertices - 3, 0, -1):
          if not matrix[v, num_vertices - 1]:
              for dst in range(v + 1, num_vertices - 1):
                  if matrix[v, dst]:
                      vertex_channels[v] = max(
                          vertex_channels[v], vertex_channels[dst]
                      )
          assert vertex_channels[v] > 0
  
      _logger.debug("vertex_channels: %s", str(vertex_channels))
  
      # Sanity check, verify that channels never increase and final channels add up.
      final_fan_in = 0
      for v in range(1, num_vertices - 1):
          if matrix[v, num_vertices - 1]:
              final_fan_in += vertex_channels[v]
          for dst in range(v + 1, num_vertices - 1):
              if matrix[v, dst]:
                  assert vertex_channels[v] >= vertex_channels[dst]
      assert final_fan_in == output_channels or num_vertices == 2
      # num_vertices == 2 means only input/output nodes, so 0 fan-in
  
      return vertex_channels
  
imports: import logging
  from collections import OrderedDict
  from typing import Callable, List, Optional, Union, Dict, Tuple, Callable, Iterable, Any, cast
  import numpy
  import torch
  import torch.nn
  from nni.mutable import Constraint, Categorical, CategoricalMultiple, Constraint, ConstraintViolation, Mutable, LabeledMutable, Sample, SampleValidationError, auto_label, label_scope
  from nni.mutable.mutable import _mutable_equal
  from nni.nas.nn.pytorch import InputChoice, LayerChoice, MutableModule
label_desc: Use image node_id 1 for calling a global function with example usage: compute_vertex_channels(input_channels, output_channels, matrix) and returns: vertex_channels, vertex_channels