Function test_sinkhorn_lpl1_vectorization

test/test_da.py:1017–1054

(nx)

@pytest.skip_backend("jax")
@pytest.skip_backend("tf")
def test_sinkhorn_lpl1_vectorization(nx):
    n_samples, n_labels = 150, 3
    rng = np.random.RandomState(42)
    M = rng.rand(n_samples, n_samples)
    labels_a = rng.randint(n_labels, size=(n_samples,))
    M, labels_a = nx.from_numpy(M), nx.from_numpy(labels_a)

    # hard-coded params from the original code
    p, epsilon = 0.5, 1e-3
    T = nx.from_numpy(rng.rand(n_samples, n_samples))

    def unvectorized(transp):
        indices_labels = []
        classes = nx.unique(labels_a)
        for c in classes:
            (idxc,) = nx.where(labels_a == c)
            indices_labels.append(idxc)
        W = nx.ones(M.shape, type_as=M)
        for i, c in enumerate(classes):
            majs = nx.sum(transp[indices_labels[i]], axis=0)
            majs = p * ((majs + epsilon) ** (p - 1))
            W[indices_labels[i]] = majs
        return W

    def vectorized(transp):
        labels_u, labels_idx = nx.unique(labels_a, return_inverse=True)
        n_labels = labels_u.shape[0]
        unroll_labels_idx = nx.eye(n_labels, type_as=transp)[labels_idx]
        W = (
            nx.repeat(transp.T[:, :, None], n_labels, axis=2)
            * unroll_labels_idx[None, :, :]
        )
        W = nx.sum(W, axis=1)
        W = p * ((W + epsilon) ** (p - 1))
        W = nx.dot(W, unroll_labels_idx.T)
        return W.T

    assert np.allclose(unvectorized(T), vectorized(T))