Support creating Bool constants in the pattern_utils (apache#7507)

Lokiiiiii · Mar 1, 2021 · 166831e · 166831e
1 parent dd1549e
commit 166831e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/relay/transforms/pattern_utils.h b/src/relay/transforms/pattern_utils.h
@@ -86,6 +86,9 @@ namespace relay {
   } else if (type == DataType::UInt(8)) {                                             \
     typedef uint8_t DType;                                                            \
     { __VA_ARGS__ }                                                                   \
+  } else if (type == DataType::Bool()) {                                              \
+    typedef bool DType;                                                               \
+    { __VA_ARGS__ }                                                                   \
   } else if ((*tvm::runtime::Registry::Get("runtime._datatype_get_type_registered"))( \
                  static_cast<uint8_t>(type.code()))) {                                \
     typedef double DType;                                                             \

diff --git a/tests/python/relay/test_pass_simplify_expr.py b/tests/python/relay/test_pass_simplify_expr.py
@@ -117,7 +117,7 @@ def after_right(x, elem_op, value):
                 assert tvm.ir.structural_equal(zz, after)
 
     for shape in [[10], [10, 10], [10, 10, 10]]:
-        for dtype in ["float32", "int32"]:
+        for dtype in ["float32", "int32", "bool"]:
             for value in [0, 1, 2]:
                 validate(shape, value, dtype)