code: def decode_mst(energy, length, has_labels=True):
      """
      Note: Counter to typical intuition, this function decodes the _maximum_
      spanning tree.
  
      Decode the optimal MST tree with the Chu-Liu-Edmonds algorithm for
      maximum spanning arboresences on graphs.
  
      Parameters
      ----------
      energy : ``numpy.ndarray``, required.
          A tensor with shape (num_labels, timesteps, timesteps)
          containing the energy of each edge. If has_labels is ``False``,
          the tensor should have shape (timesteps, timesteps) instead.
      length : ``int``, required.
          The length of this sequence, as the energy may have come
          from a padded batch.
      has_labels : ``bool``, optional, (default = True)
          Whether the graph has labels or not.
      """
      if has_labels and energy.ndim != 3:
          raise ConfigurationError(
              "The dimension of the energy array is not equal to 3."
          )
      elif not has_labels and energy.ndim != 2:
          raise ConfigurationError(
              "The dimension of the energy array is not equal to 2."
          )
      input_shape = energy.shape
      max_length = input_shape[-1]
  
      # Our energy matrix might have been batched -
      # here we clip it to contain only non padded tokens.
      if has_labels:
          energy = energy[:, :length, :length]
          # get best label for each edge.
          label_id_matrix = energy.argmax(axis=0)
          energy = energy.max(axis=0)
      else:
          energy = energy[:length, :length]
          label_id_matrix = None
      # get original score matrix
      original_score_matrix = energy
      # initialize score matrix to original score matrix
      score_matrix = numpy.array(original_score_matrix, copy=True)
  
      old_input = numpy.zeros([length, length], dtype=numpy.int32)
      old_output = numpy.zeros([length, length], dtype=numpy.int32)
      current_nodes = [True for _ in range(length)]
      representatives = []
  
      for node1 in range(length):
          original_score_matrix[node1, node1] = 0.0
          score_matrix[node1, node1] = 0.0
          representatives.append(set([node1]))
  
          for node2 in range(node1 + 1, length):
              old_input[node1, node2] = node1
              old_output[node1, node2] = node2
  
              old_input[node2, node1] = node2
              old_output[node2, node1] = node1
  
      final_edges = {}
  
      # The main algorithm operates inplace.
      chu_liu_edmonds(
          length,
          score_matrix,
          current_nodes,
          final_edges,
          old_input,
          old_output,
          representatives,
      )
  
      heads = numpy.zeros([max_length], numpy.int32)
      if has_labels:
          head_type = numpy.ones([max_length], numpy.int32)
      else:
          head_type = None
  
      for child, parent in list(final_edges.items()):
          heads[child] = parent
          if has_labels:
              head_type[child] = label_id_matrix[parent, child]
  
      return heads, head_type
  
imports: from __future__ import absolute_import
  import numpy
  from allennlp.common.checks import ConfigurationError
label_desc: Use image node_id 1 for calling a global function with example usage: decode_mst(energy, length, has_labels) and returns: heads,  head_type