Function test_stack

imperative/python/test/unit/functional/test_tensor.py:176–217

(is_varnode)

@pytest.mark.parametrize("is_varnode", [True, False])
def test_stack(is_varnode):
    if is_varnode:
        network = Network()
    else:
        network = None

    data1 = np.random.random((3, 2, 2)).astype("float32")
    data2 = np.random.random((3, 2, 2)).astype("float32")
    data3 = np.random.random((3, 2, 2)).astype("float32")

    cases = [{"input": [data1, data2]}, {"input": [data1, data3]}]
    for ai in range(3):

        def run(data1, data2):
            return F.stack([data1, data2], axis=ai)

        opr_test(
            cases, run, ref_fn=lambda x, y: np.stack([x, y], axis=ai), network=network
        )

    x1 = Tensor(np.arange(0, 3, dtype=np.float32).reshape((3)))
    x2 = Tensor(np.arange(6, 9, dtype=np.float32).reshape((3)))
    y = F.stack([x1, x2], axis=-1)
    np.testing.assert_equal(
        y.numpy(), np.array([[0, 6], [1, 7], [2, 8]]).astype(np.float32)
    )

    x1 = Tensor(np.arange(0, 3, dtype=np.float32).reshape((3)))
    x2 = Tensor(np.arange(6, 9, dtype=np.float32).reshape((3)))
    y = F.stack([x1, x2], axis=-1)
    np.testing.assert_equal(
        y.numpy(), np.array([[0, 6], [1, 7], [2, 8]]).astype(np.float32)
    )

    x1 = Tensor(np.random.rand(600))
    x2 = F.broadcast_to(Tensor(np.array(3)), (600,))

    y = F.stack([x2, x1], axis=0)
    np.testing.assert_equal(y.numpy(), np.stack((x2.numpy(), x1.numpy()), axis=0))

    y = F.stack([x2, x2], axis=0)
    np.testing.assert_equal(y.numpy(), np.stack((x2.numpy(), x2.numpy()), axis=0))


@pytest.mark.parametrize("is_varnode", [True, False])