code: def _mst(scores):
      r"""
      with some modification to support parser output for MST decoding
      https://github.com/tdozat/Parser/blob/0739216129cd39d69997d28cbc4133b360ea3934/lib/models/nn.py#L692
      """
      length = scores.shape[0]
      min_score = scores.min() - 1
      eye = np.eye(length)
      scores = scores * (1 - eye) + min_score * eye
      heads = np.argmax(scores, axis=1)
      heads[0] = 0
      tokens = np.arange(1, length)
      roots = np.where(heads[tokens] == 0)[0] + 1
      if len(roots) < 1:
          root_scores = scores[tokens, 0]
          head_scores = scores[tokens, heads[tokens]]
          new_root = tokens[np.argmax(root_scores / head_scores)]
          heads[new_root] = 0
      elif len(roots) > 1:
          root_scores = scores[roots, 0]
          scores[roots, 0] = 0
          new_heads = np.argmax(scores[roots][:, tokens], axis=1) + 1
          new_root = roots[
              np.argmin(scores[roots, new_heads] / root_scores)
          ]
          heads[roots] = new_heads
          heads[new_root] = 0
  
      edges = defaultdict(set)
      vertices = set((0,))
      for dep, head in enumerate(heads[tokens]):
          vertices.add(dep + 1)
          edges[head].add(dep + 1)
      for cycle in _find_cycle(vertices, edges):
          dependents = set()
          to_visit = set(cycle)
          while len(to_visit) > 0:
              node = to_visit.pop()
              if node not in dependents:
                  dependents.add(node)
                  to_visit.update(edges[node])
          cycle = np.array(list(cycle))
          old_heads = heads[cycle]
          old_scores = scores[cycle, old_heads]
          non_heads = np.array(list(dependents))
          scores[
              np.repeat(cycle, len(non_heads)),
              np.repeat([non_heads], len(cycle), axis=0).flatten(),
          ] = min_score
          new_heads = np.argmax(scores[cycle][:, tokens], axis=1) + 1
          new_scores = scores[cycle, new_heads] / old_scores
          change = np.argmax(new_scores)
          changed_cycle = cycle[change]
          old_head = old_heads[change]
          new_head = new_heads[change]
          heads[changed_cycle] = new_head
          edges[new_head].add(changed_cycle)
          edges[old_head].remove(changed_cycle)
  
      return heads
  
imports: from collections import defaultdict
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from ...core.utils import seq_len_to_mask
  from ...embeddings.torch.utils import get_embeddings
  from ...modules.torch.dropout import TimestepDropout
  from ...modules.torch.encoder.transformer import TransformerEncoder
  from ...modules.torch.encoder.variational_rnn import VarLSTM
label_desc: Use image node_id 1 for calling a global function with example usage: _mst(scores) and returns: heads