code: def mst(scores, mask, multiroot=False):
      r"""
      MST algorithm for decoding non-pojective trees.
      This is a wrapper for ChuLiu/Edmonds algorithm.
  
      The algorithm first runs ChuLiu/Edmonds to parse a tree and then have a check of multi-roots,
      If ``multiroot=True`` and there indeed exist multi-roots, the algorithm seeks to find
      best single-root trees by iterating all possible single-root trees parsed by ChuLiu/Edmonds.
      Otherwise the resulting trees are directly taken as the final outputs.
  
      Args:
          scores (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.
              Scores of all dependent-head pairs.
          mask (~torch.BoolTensor): ``[batch_size, seq_len]``.
              The mask to avoid parsing over padding tokens.
              The first column serving as pseudo words for roots should be ``False``.
          muliroot (bool):
              Ensures to parse a single-root tree If ``False``.
  
      Returns:
          ~torch.Tensor:
              A tensor with shape ``[batch_size, seq_len]`` for the resulting non-projective parse trees.
  
      Examples:
          >>> scores = torch.tensor([[[-11.9436, -13.1464,  -6.4789, -13.8917],
                                      [-60.6957, -60.2866, -48.6457, -63.8125],
                                      [-38.1747, -49.9296, -45.2733, -49.5571],
                                      [-19.7504, -23.9066,  -9.9139, -16.2088]]])
          >>> scores[:, 0, 1:] = float('-inf')
          >>> scores.diagonal(0, 1, 2)[1:].fill_(float('-inf'))
          >>> mask = torch.tensor([[False,  True,  True,  True]])
          >>> mst(scores, mask)
          tensor([[0, 2, 0, 2]])
      """
  
      batch_size, seq_len, _ = scores.shape
      scores = scores.cpu().unbind()
  
      preds = []
      for i, length in enumerate(mask.sum(1).tolist()):
          s = scores[i][: length + 1, : length + 1]
          tree = chuliu_edmonds(s)
          roots = torch.where(tree[1:].eq(0))[0] + 1
          if not multiroot and len(roots) > 1:
              s_root = s[:, 0]
              s_best = float("-inf")
              s = s.index_fill(1, torch.tensor(0), float("-inf"))
              for root in roots:
                  s[:, 0] = float("-inf")
                  s[root, 0] = s_root[root]
                  t = chuliu_edmonds(s)
                  s_tree = s[1:].gather(1, t[1:].unsqueeze(-1)).sum()
                  if s_tree > s_best:
                      s_best, tree = s_tree, t
          preds.append(tree)
  
      return pad(preds, total_length=seq_len).to(mask.device)
  
file_path: underthesea/underthesea/utils/sp_alg.py
imports: import torch
  from underthesea.utils.sp_fn import pad, stripe
label_desc: Use image node_id 4 for calling a global function with example usage: mst(scores, mask, multiroot) and returns: pad