Function test_LazyTensor

test/test_utils.py:499–536

(nx)

def test_LazyTensor(nx):
    n1 = 100
    n2 = 200
    shape = (n1, n2)

    rng = np.random.RandomState(42)
    x1 = rng.randn(n1, 2)
    x2 = rng.randn(n2, 2)

    x1, x2 = nx.from_numpy(x1, x2)

    # i,j can be integers or slices, x1,x2 have to be passed as keyword arguments
    def getitem(i, j, x1, x2):
        return nx.dot(x1[i], x2[j].T)

    # create a lazy tensor
    T = ot.utils.LazyTensor((n1, n2), getitem, x1=x1, x2=x2)

    assert T.shape == (n1, n2)
    assert str(T) == "LazyTensor(shape=(100, 200),attributes=(x1,x2))"

    assert T.x1 is x1
    assert T.x2 is x2

    # get the full tensor (not lazy)
    assert T[:].shape == shape

    # get one component
    assert T[1, 1] == nx.dot(x1[1], x2[1].T)

    # get one row
    assert T[1].shape == (n2,)

    # get one column with slices
    assert T[::10, 5].shape == (10,)

    with pytest.raises(NotImplementedError):
        T["error"]


def test_OTResult_LazyTensor(nx):