Function test_LazyTensor_reduce

test/test_utils.py:548–596

(nx)

def test_LazyTensor_reduce(nx):
    T, a, b = get_LazyTensor(nx)

    T0 = T[:]
    s0 = nx.sum(T0)

    # total sum
    s = ot.utils.reduce_lazytensor(T, nx.sum, nx=nx)
    np.testing.assert_allclose(nx.to_numpy(s), 1)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(s0))

    s2 = ot.utils.reduce_lazytensor(T, nx.sum)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(s2))

    s2 = ot.utils.reduce_lazytensor(T, nx.sum, batch_size=500)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(s2))

    s2 = ot.utils.reduce_lazytensor(T, nx.sum, batch_size=11)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(s2))

    # sum over axis 0
    s = ot.utils.reduce_lazytensor(T, nx.sum, axis=0, nx=nx)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(b))

    # sum over axis 1
    s = ot.utils.reduce_lazytensor(T, nx.sum, axis=1, nx=nx)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(a))

    # test otehr reduction function
    s = ot.utils.reduce_lazytensor(T, nx.logsumexp, axis=1, nx=nx)
    s2 = nx.logsumexp(T[:], axis=1)
    np.testing.assert_allclose(nx.to_numpy(s), nx.to_numpy(s2))

    # test 3D tensors
    def getitem(i, j, k, a, b, c):
        return a[i, None, None] * b[None, j, None] * c[None, None, k]

    # create a lazy tensor
    n = a.shape[0]
    T = ot.utils.LazyTensor((n, n, n), getitem, a=a, b=a, c=a)

    # total sum
    s1 = ot.utils.reduce_lazytensor(T, nx.sum, axis=0, nx=nx)
    s2 = ot.utils.reduce_lazytensor(T, nx.sum, axis=1, nx=nx)

    np.testing.assert_allclose(nx.to_numpy(s1), nx.to_numpy(s2))

    with pytest.raises(NotImplementedError):
        ot.utils.reduce_lazytensor(T, nx.sum, axis=2, nx=nx, batch_size=10)


def test_lowrank_LazyTensor(nx):