[RUNTIME][FFI] Fix cython FFI compact with np.int64 (apache#6321)

trevor-m · Aug 26, 2020 · d87c961 · d87c961
1 parent a0cdb57
commit d87c961
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/python/tvm/_ffi/_cython/packed_func.pxi b/python/tvm/_ffi/_cython/packed_func.pxi
@@ -116,7 +116,7 @@ cdef inline int make_arg(object arg,
         ptr = arg._tvm_handle
         value[0].v_handle = (<void*>ptr)
         tcode[0] = arg.__class__._tvm_tcode
-    elif isinstance(arg, (int, long)):
+    elif isinstance(arg, Integral):
         value[0].v_int64 = arg
         tcode[0] = kInt
     elif isinstance(arg, float):

diff --git a/tests/python/unittest/test_runtime_packed_func.py b/tests/python/unittest/test_runtime_packed_func.py
@@ -296,11 +296,14 @@ def check_assign(dtype):
     for t in ["float64", "float32"]:
         check_assign(t)
 
+def test_numpy_scalar():
+    maxint = (1<<63) - 1
+    assert tvm.testing.echo(np.int64(maxint)) == maxint
 
 
 if __name__ == "__main__":
+    test_numpy_scalar()
     test_rvalue_ref()
-    exit(0)
     test_empty_array()
     test_get_global()
     test_get_callback_with_node()