Fix boolean dtype error in hybridization

src/nnvm/plan_memory.cc

@@ -38,21 +38,22 @@ namespace {
 // Return bytes of data flag.
 static int MXGetDTypeSize(int type_flag) {
   switch (type_flag) {
-    case kUint8:
-    case kInt8:
+    case mshadow::kUint8:
+    case mshadow::kInt8:
+    case mshadow::kBool:
       return 1;
-    case kFloat16:
-    case kBfloat16:
-    case kInt16:
-    case kUint16:
+    case mshadow::kFloat16:
+    case mshadow::kBfloat16:
+    case mshadow::kInt16:
+    case mshadow::kUint16:
       return 2;
-    case kFloat32:
-    case kInt32:
-    case kUint32:
+    case mshadow::kFloat32:
+    case mshadow::kInt32:
+    case mshadow::kUint32:
       return 4;
-    case kFloat64:
-    case kInt64:
-    case kUint64:
+    case mshadow::kFloat64:
+    case mshadow::kInt64:
+    case mshadow::kUint64:
       return 8;
     default:
       LOG(FATAL) << "unknown type_flag=" << type_flag;

tests/python/unittest/test_numpy_gluon.py

@@ -400,6 +400,27 @@ def hybrid_forward(self, F, x, y):
                                   mx.np.random.normal(0, 1, (10, 5, 8))])
 
 
+@with_seed()
+@use_np
+def test_hybridize_boolean_dtype():
+    class Foo(gluon.HybridBlock):
+        def __init__(self, prefix=None, params=None):
+            super(Foo, self).__init__(prefix=prefix, params=params)
+
+        def hybrid_forward(self, F, valid_length):
+            mask = ((F.np.ones((10,)) / 2) < valid_length)
+            return mask
+
+    valid_length = mx.np.random.uniform(size=(10,))
+    foo = Foo()
+    out1 = foo(valid_length)
+
+    foo = Foo()
+    foo.hybridize()
+    out2 = foo(valid_length)
+
+    assert mx.test_utils.same(out1.asnumpy(), out2.asnumpy())
+
 
 if __name__ == '__main__':
     import nose