Function test_condtake

imperative/python/test/unit/functional/test_tensor.py:112–128

(is_varnode)

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

    x = np.array([[1, 2, 3], [4, 5, 6]]).astype("float32")
    y = np.array([[True, False, True], [False, True, True]])
    xx = make_tensor(x, network)
    yy = make_tensor(y, network)
    val, idx = F.cond_take(yy, xx)
    if is_varnode:
        np.testing.assert_equal(get_var_value(val), x[y])
        np.testing.assert_equal(get_var_value(idx), np.where(y.reshape(-1))[0])
    else:
        np.testing.assert_equal(val.numpy(), x[y])
        np.testing.assert_equal(idx.numpy(), np.where(y.reshape(-1))[0])


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