diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc index 22844745982fd..07cae5e2c746f 100644 --- a/src/tir/ir/expr.cc +++ b/src/tir/ir/expr.cc @@ -158,7 +158,8 @@ PrimExpr SelectNode::make(PrimExpr condition, PrimExpr true_value, PrimExpr fals CHECK(true_value.defined()) << "ValueError: true_value is undefined"; CHECK(false_value.defined()) << "ValueError: true_value is undefined"; CHECK(condition.dtype().is_bool()); - CHECK_EQ(condition.dtype().lanes(), true_value.dtype().lanes()); + CHECK(condition.dtype().lanes() == true_value.dtype().lanes() || + condition.dtype().lanes() == 1); CHECK(false_value.dtype() == true_value.dtype()) << "TypeError: mismatched types"; ObjectPtr node = make_object(); diff --git a/tests/python/relay/test_op_level4.py b/tests/python/relay/test_op_level4.py index c5cd708187953..44b51f2c23676 100644 --- a/tests/python/relay/test_op_level4.py +++ b/tests/python/relay/test_op_level4.py @@ -323,6 +323,7 @@ def verify(dshape, begin, end, strides, vshape, test_ref=True): op_res = intrp.evaluate(func)(x_data, v_data) tvm.testing.assert_allclose(op_res.asnumpy(), ref_res) + verify((3, 4, 16), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2)) verify((3, 4, 3), [0, 0, 0], [4, -5, 4], [1, -1, 2], (3, 1, 2)) verify((3, 4, 3), [1, 1, 0], [4, 4, 3], [2, 1, 1], (1, 3, 3)) verify((3, 4, 3), [1, -1, 0], [4, -5, 3], [2, -1, 1], (1, 4, 3))