Function test_squeeze

imperative/python/test/unit/functional/test_tensor.py:478–498

(is_varnode)

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

    x = Tensor(np.array([1, 2], dtype=np.int32).reshape(1, 1, 2, 1))
    y = F.squeeze(x, -1)
    np.testing.assert_equal(y.numpy(), np.array([[[1, 2]]]).astype(np.int32))

    x = np.arange(6, dtype="float32").reshape(1, 2, 3, 1)
    xx = make_tensor(x, network)

    for axis in [None, 3, -4, (3, -4)]:
        y = np.squeeze(x, axis)
        yy = F.squeeze(xx, axis)
        np.testing.assert_equal(y, yy.numpy())

    if is_varnode:
        set_symbolic_shape(saved_symbolic_shape)


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