Numpy compatible dtype inference for tvm.convert and tvm.const

python/tvm/_ffi/node_generic.py

@@ -25,11 +25,29 @@
 # Node base class
 _CLASS_NODE_BASE = None
 
+
 def _set_class_node_base(cls):
     global _CLASS_NODE_BASE
     _CLASS_NODE_BASE = cls
 
 
+def _scalar_type_inference(value):
+    if hasattr(value, 'dtype'):
+        dtype = str(value.dtype)
+    elif isinstance(value, bool):
+        dtype = 'bool'
+    elif isinstance(value, float):
+        # We intentionally convert the float to float32 since it's more common in DL.
+        dtype = 'float32'
+    elif isinstance(value, int):
+        # We intentionally convert the python int to int32 since it's more common in DL.
+        dtype = 'int32'
+    else:
+        raise NotImplementedError('Cannot automatically inference the type.'
+                                  ' value={}'.format(value))
+    return dtype
+
+
 class NodeGeneric(object):
     """Base class for all classes that can be converted to node."""
     def asnode(self):
@@ -86,7 +104,7 @@ def const(value, dtype=None):
     value : int or float
         The input value
 
-    dtype : str
+    dtype : str or None, optional
         The data type.
 
     Returns
@@ -95,8 +113,5 @@ def const(value, dtype=None):
         Constant expression corresponds to the value.
     """
     if dtype is None:
-        if isinstance(value, Integral):
-            dtype = 'int32'
-        else:
-            dtype = 'float32'
+        dtype = _scalar_type_inference(value)
     return _api_internal._const(value, dtype)

python/tvm/api.py

@@ -23,6 +23,7 @@
 from ._ffi.base import string_types
 from ._ffi.node import register_node, NodeBase
 from ._ffi.node import convert_to_node as _convert_to_node
+from ._ffi.node_generic import _scalar_type_inference
 from ._ffi.function import Function
 from ._ffi.function import _init_api, register_func, get_global_func, extract_ext_funcs
 from ._ffi.function import convert_to_tvm_func as _convert_tvm_func
@@ -73,22 +74,24 @@ def max_value(dtype):
     return _api_internal._max_value(dtype)
 
 
-def const(value, dtype):
+def const(value, dtype=None):
     """construct a constant
 
     Parameters
     ----------
     value : number
         The content of the constant number.
 
-    dtype : str
+    dtype : str or None, optional
         The data type.
 
     Returns
     -------
     const_val: tvm.Expr
         The result expression.
     """
+    if dtype is None:
+        dtype = _scalar_type_inference(value)
     return _api_internal._const(value, dtype)
 
 

tests/python/unittest/test_lang_basic.py

@@ -15,13 +15,30 @@
 # specific language governing permissions and limitations
 # under the License.
 import tvm
+import numpy as np
 
 def test_const():
     x = tvm.const(1, "int32")
     print(x.dtype)
     assert x.dtype == tvm.int32
     assert isinstance(x, tvm.expr.IntImm)
 
+
+def test_scalar_dtype_inference():
+    for data in [True, np.bool(1), np.uint8(1), np.uint16(1), np.uint32(1), np.uint64(1),
+                 np.int8(1), np.int16(1), np.int32(1), np.int64(1),
+                 np.float16(1), np.float32(1), np.float64(1)]:
+        assert tvm.const(data).dtype == str(np.array(data).dtype)
+    assert tvm.const(1).dtype == 'int32'
+    assert tvm.const(1.0).dtype == 'float32'
+
+    for data in [True, np.bool(1), np.uint8(1), np.uint16(1), np.uint32(1), np.uint64(1),
+                 np.int8(1), np.int16(1), np.int32(1), np.int64(1),
+                 np.float16(1), np.float32(1), np.float64(1)]:
+        assert tvm.convert(data).dtype == str(np.array(data).dtype)
+    assert tvm.convert(1).dtype == 'int32'
+    assert tvm.convert(1.0).dtype == 'float32'
+
 def test_make():
     x = tvm.const(1, "int32")
     y = tvm.var("x")
@@ -175,6 +192,7 @@ def test_equality_string_imm():
     test_cast()
     test_attr()
     test_const()
+    test_scalar_dtype_inference()
     test_make()
     test_ir()
     test_basic()