code: def chuliu_edmonds(s):
      r"""
      ChuLiu/Edmonds algorithm for non-projective decoding.
  
      Some code is borrowed from `tdozat's implementation`_.
      Descriptions of notations and formulas can be found in
      `Non-projective Dependency Parsing using Spanning Tree Algorithms`_.
  
      Notes:
          The algorithm does not guarantee to parse a single-root tree.
  
      References:
          - Ryan McDonald, Fernando Pereira, Kiril Ribarov and Jan Hajic. 2005.
            `Non-projective Dependency Parsing using Spanning Tree Algorithms`_.
  
      Args:
          s (~torch.Tensor): ``[seq_len, seq_len]``.
              Scores of all dependent-head pairs.
  
      Returns:
          ~torch.Tensor:
              A tensor with shape ``[seq_len]`` for the resulting non-projective parse tree.
  
      .. _tdozat's implementation:
          https://github.com/tdozat/Parser-v3
      .. _Non-projective Dependency Parsing using Spanning Tree Algorithms:
          https://www.aclweb.org/anthology/H05-1066/
      """
  
      s[0, 1:] = float("-inf")
      # prevent self-loops
      s.diagonal()[1:].fill_(float("-inf"))
      # select heads with highest scores
      tree = s.argmax(-1)
      # return the cycle finded by tarjan algorithm lazily
      cycle = next(tarjan(tree.tolist()[1:]), None)
      # if the tree has no cycles, then it is a MST
      if not cycle:
          return tree
      # indices of cycle in the original tree
      cycle = torch.tensor(cycle)
      # indices of noncycle in the original tree
      noncycle = torch.ones(len(s)).index_fill_(0, cycle, 0)
      noncycle = torch.where(noncycle.gt(0))[0]
  
      def contract(s):
          # heads of cycle in original tree
          cycle_heads = tree[cycle]
          # scores of cycle in original tree
          s_cycle = s[cycle, cycle_heads]
  
          # calculate the scores of cycle's potential dependents
          # s(c->x) = max(s(x'->x)), x in noncycle and x' in cycle
          s_dep = s[noncycle][:, cycle]
          # find the best cycle head for each noncycle dependent
          deps = s_dep.argmax(1)
          # calculate the scores of cycle's potential heads
          # s(x->c) = max(s(x'->x) - s(a(x')->x') + s(cycle)), x in noncycle and x' in cycle
          #                                                    a(v) is the predecessor of v in cycle
          #                                                    s(cycle) = sum(s(a(v)->v))
          s_head = (
              s[cycle][:, noncycle]
              - s_cycle.view(-1, 1)
              + s_cycle.sum()
          )
          # find the best noncycle head for each cycle dependent
          heads = s_head.argmax(0)
  
          contracted = torch.cat((noncycle, torch.tensor([-1])))
          # calculate the scores of contracted graph
          s = s[contracted][:, contracted]
          # set the contracted graph scores of cycle's potential dependents
          s[:-1, -1] = s_dep[range(len(deps)), deps]
          # set the contracted graph scores of cycle's potential heads
          s[-1, :-1] = s_head[heads, range(len(heads))]
  
          return s, heads, deps
  
      # keep track of the endpoints of the edges into and out of cycle for reconstruction later
      s, heads, deps = contract(s)
  
      # y is the contracted tree
      y = chuliu_edmonds(s)
      # exclude head of cycle from y
      y, cycle_head = y[:-1], y[-1]
  
      # fix the subtree with no heads coming from the cycle
      # len(y) denotes heads coming from the cycle
      subtree = y < len(y)
      # add the nodes to the new tree
      tree[noncycle[subtree]] = noncycle[y[subtree]]
      # fix the subtree with heads coming from the cycle
      subtree = ~subtree
      # add the nodes to the tree
      tree[noncycle[subtree]] = cycle[deps[subtree]]
      # fix the root of the cycle
      cycle_root = heads[cycle_head]
      # break the cycle and add the root of the cycle to the tree
      tree[cycle[cycle_root]] = noncycle[cycle_head]
  
      return tree
  
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 3 for calling a global function with example usage: chuliu_edmonds(s) and returns: tree, tree, s,  heads,  deps