def chuliu_edmonds(s):
    "\n    ChuLiu/Edmonds algorithm for non-projective decoding :cite:`mcdonald-etal-2005-non`.\n\n    Some code is borrowed from `tdozat's implementation`_.\n    Descriptions of notations and formulas can be found in :cite:`mcdonald-etal-2005-non`.\n\n    Notes:\n        The algorithm does not guarantee to parse a single-root tree.\n\n    Args:\n        s (~torch.Tensor): ``[seq_len, seq_len]``.\n            Scores of all dependent-head pairs.\n\n    Returns:\n        ~torch.Tensor:\n            A tensor with shape ``[seq_len]`` for the resulting non-projective parse tree.\n\n    .. _tdozat's implementation:\n        https://github.com/tdozat/Parser-v3\n    "
    s[0, 1:] = MIN
    s.diagonal()[1:].fill_(MIN)
    tree = s.argmax((- 1))
    cycle = next(tarjan(tree.tolist()[1:]), None)
    if (not cycle):
        return tree
    cycle = torch.tensor(cycle)
    noncycle = torch.ones(len(s)).index_fill_(0, cycle, 0)
    noncycle = torch.where(noncycle.gt(0))[0]

    def contract(s):
        cycle_heads = tree[cycle]
        s_cycle = s[(cycle, cycle_heads)]
        s_dep = s[noncycle][:, cycle]
        deps = s_dep.argmax(1)
        s_head = ((s[cycle][:, noncycle] - s_cycle.view((- 1), 1)) + s_cycle.sum())
        heads = s_head.argmax(0)
        contracted = torch.cat((noncycle, torch.tensor([(- 1)])))
        s = s[contracted][:, contracted]
        s[:(- 1), (- 1)] = s_dep[(range(len(deps)), deps)]
        s[(- 1), :(- 1)] = s_head[(heads, range(len(heads)))]
        return (s, heads, deps)
    (s, heads, deps) = contract(s)
    y = chuliu_edmonds(s)
    (y, cycle_head) = (y[:(- 1)], y[(- 1)])
    subtree = (y < len(y))
    tree[noncycle[subtree]] = noncycle[y[subtree]]
    subtree = (~ subtree)
    tree[noncycle[subtree]] = cycle[deps[subtree]]
    cycle_root = heads[cycle_head]
    tree[cycle[cycle_root]] = noncycle[cycle_head]
    return tree
ChuLiu/Edmonds algorithm for non-projective decoding :cite:`mcdonald-etal-2005-non`.

Some code is borrowed from `tdozat's implementation`_.
Descriptions of notations and formulas can be found in :cite:`mcdonald-etal-2005-non`.

Notes:
    The algorithm does not guarantee to parse a single-root tree.

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
def chuliu_edmonds(s):
    "\n    ChuLiu/Edmonds algorithm for non-projective decoding :cite:`mcdonald-etal-2005-non`.\n\n    Some code is borrowed from `tdozat's implementation`_.\n    Descriptions of notations and formulas can be found in :cite:`mcdonald-etal-2005-non`.\n\n    Notes:\n        The algorithm does not guarantee to parse a single-root tree.\n\n    Args:\n        s (~torch.Tensor): ``[seq_len, seq_len]``.\n            Scores of all dependent-head pairs.\n\n    Returns:\n        ~torch.Tensor:\n            A tensor with shape ``[seq_len]`` for the resulting non-projective parse tree.\n\n    .. _tdozat's implementation:\n        https://github.com/tdozat/Parser-v3\n    "
    s[0, 1:] = MIN
    s.diagonal()[1:].fill_(MIN)
    tree = s.argmax((- 1))
    cycle = next(tarjan(tree.tolist()[1:]), None)
    if (not cycle):
        return tree
    cycle = torch.tensor(cycle)
    noncycle = torch.ones(len(s)).index_fill_(0, cycle, 0)
    noncycle = torch.where(noncycle.gt(0))[0]

    def contract(s):
        cycle_heads = tree[cycle]
        s_cycle = s[(cycle, cycle_heads)]
        s_dep = s[noncycle][:, cycle]
        deps = s_dep.argmax(1)
        s_head = ((s[cycle][:, noncycle] - s_cycle.view((- 1), 1)) + s_cycle.sum())
        heads = s_head.argmax(0)
        contracted = torch.cat((noncycle, torch.tensor([(- 1)])))
        s = s[contracted][:, contracted]
        s[:(- 1), (- 1)] = s_dep[(range(len(deps)), deps)]
        s[(- 1), :(- 1)] = s_head[(heads, range(len(heads)))]
        return (s, heads, deps)
    (s, heads, deps) = contract(s)
    y = chuliu_edmonds(s)
    (y, cycle_head) = (y[:(- 1)], y[(- 1)])
    subtree = (y < len(y))
    tree[noncycle[subtree]] = noncycle[y[subtree]]
    subtree = (~ subtree)
    tree[noncycle[subtree]] = cycle[deps[subtree]]
    cycle_root = heads[cycle_head]
    tree[cycle[cycle_root]] = noncycle[cycle_head]
    return tree