def mst(scores, mask, multiroot=False):
    '\n    MST algorithm for decoding non-projective trees.\n    This is a wrapper for ChuLiu/Edmonds algorithm.\n\n    The algorithm first runs ChuLiu/Edmonds to parse a tree and then have a check of multi-roots,\n    If ``multiroot=True`` and there indeed exist multi-roots, the algorithm seeks to find\n    best single-root trees by iterating all possible single-root trees parsed by ChuLiu/Edmonds.\n    Otherwise the resulting trees are directly taken as the final outputs.\n\n    Args:\n        scores (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.\n            Scores of all dependent-head pairs.\n        mask (~torch.BoolTensor): ``[batch_size, seq_len]``.\n            The mask to avoid parsing over padding tokens.\n            The first column serving as pseudo words for roots should be ``False``.\n        multiroot (bool):\n            Ensures to parse a single-root tree If ``False``.\n\n    Returns:\n        ~torch.Tensor:\n            A tensor with shape ``[batch_size, seq_len]`` for the resulting non-projective parse trees.\n\n    Examples:\n        >>> scores = torch.tensor([[[-11.9436, -13.1464,  -6.4789, -13.8917],\n                                    [-60.6957, -60.2866, -48.6457, -63.8125],\n                                    [-38.1747, -49.9296, -45.2733, -49.5571],\n                                    [-19.7504, -23.9066,  -9.9139, -16.2088]]])\n        >>> scores[:, 0, 1:] = MIN\n        >>> scores.diagonal(0, 1, 2)[1:].fill_(MIN)\n        >>> mask = torch.tensor([[False,  True,  True,  True]])\n        >>> mst(scores, mask)\n        tensor([[0, 2, 0, 2]])\n    '
    (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 = MIN
            s = s.index_fill(1, torch.tensor(0), MIN)
            for root in roots:
                s[:, 0] = MIN
                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)
MST algorithm for decoding non-projective 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``.
    multiroot (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:] = MIN
    >>> scores.diagonal(0, 1, 2)[1:].fill_(MIN)
    >>> mask = torch.tensor([[False,  True,  True,  True]])
    >>> mst(scores, mask)
    tensor([[0, 2, 0, 2]])
def mst(scores, mask, multiroot=False):
    '\n    MST algorithm for decoding non-projective trees.\n    This is a wrapper for ChuLiu/Edmonds algorithm.\n\n    The algorithm first runs ChuLiu/Edmonds to parse a tree and then have a check of multi-roots,\n    If ``multiroot=True`` and there indeed exist multi-roots, the algorithm seeks to find\n    best single-root trees by iterating all possible single-root trees parsed by ChuLiu/Edmonds.\n    Otherwise the resulting trees are directly taken as the final outputs.\n\n    Args:\n        scores (~torch.Tensor): ``[batch_size, seq_len, seq_len]``.\n            Scores of all dependent-head pairs.\n        mask (~torch.BoolTensor): ``[batch_size, seq_len]``.\n            The mask to avoid parsing over padding tokens.\n            The first column serving as pseudo words for roots should be ``False``.\n        multiroot (bool):\n            Ensures to parse a single-root tree If ``False``.\n\n    Returns:\n        ~torch.Tensor:\n            A tensor with shape ``[batch_size, seq_len]`` for the resulting non-projective parse trees.\n\n    Examples:\n        >>> scores = torch.tensor([[[-11.9436, -13.1464,  -6.4789, -13.8917],\n                                    [-60.6957, -60.2866, -48.6457, -63.8125],\n                                    [-38.1747, -49.9296, -45.2733, -49.5571],\n                                    [-19.7504, -23.9066,  -9.9139, -16.2088]]])\n        >>> scores[:, 0, 1:] = MIN\n        >>> scores.diagonal(0, 1, 2)[1:].fill_(MIN)\n        >>> mask = torch.tensor([[False,  True,  True,  True]])\n        >>> mst(scores, mask)\n        tensor([[0, 2, 0, 2]])\n    '
    (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 = MIN
            s = s.index_fill(1, torch.tensor(0), MIN)
            for root in roots:
                s[:, 0] = MIN
                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)