code: def prune(
      matrix, ops
  ) -> Tuple[np.ndarray, List[Union[str, Callable[[int], nn.Module]]]]:
      """
      Prune the extraneous parts of the graph.
  
      General procedure:
  
      1. Remove parts of graph not connected to input.
      2. Remove parts of graph not connected to output.
      3. Reorder the vertices so that they are consecutive after steps 1 and 2.
  
      These 3 steps can be combined by deleting the rows and columns of the
      vertices that are not reachable from both the input and output (in reverse).
      """
      num_vertices = np.shape(matrix)[0]
  
      # calculate the connection matrix within V number of steps.
      connections = np.linalg.matrix_power(
          matrix + np.eye(num_vertices), num_vertices
      )
  
      visited_from_input = set(
          [i for i in range(num_vertices) if connections[0, i]]
      )
      visited_from_output = set(
          [i for i in range(num_vertices) if connections[i, -1]]
      )
  
      # Any vertex that isn't connected to both input and output is extraneous to the computation graph.
      extraneous = set(range(num_vertices)).difference(
          visited_from_input.intersection(visited_from_output)
      )
  
      if len(extraneous) > num_vertices - 2:
          raise ConstraintViolation(
              "Non-extraneous graph is less than 2 vertices, "
              "the input is not connected to the output and the spec is invalid."
          )
  
      matrix = np.delete(matrix, list(extraneous), axis=0)
      matrix = np.delete(matrix, list(extraneous), axis=1)
      for index in sorted(extraneous, reverse=True):
          del ops[index]
      return matrix, ops
  
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 2 for calling a global function with example usage: prune(matrix, ops) and returns: matrix,  ops