Function f2

test/test_da.py:1001–1003

(G)

    unroll_labels_idx = nx.eye(n_labels, type_as=labels_u)[None, labels_idx]

    def f2(G):
        G_split = nx.repeat(G.T[:, :, None], n_labels, axis=2)
        return nx.sum(nx.norm(G_split * unroll_labels_idx, axis=1))

    def df2(G):
        G_split = nx.repeat(G.T[:, :, None], n_labels, axis=2) * unroll_labels_idx