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)
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)