[bug] Accept numpy integers in ndrange (#5245)

taichi-dev · Jul 4, 2022 · 42a9cae · 42a9cae
1 parent 752a843
commit 42a9cae
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/python/taichi/lang/_ndrange.py b/python/taichi/lang/_ndrange.py
@@ -1,5 +1,6 @@
 import collections.abc
 
+import numpy as np
 from taichi.lang import ops
 from taichi.lang.exception import TaichiSyntaxError, TaichiTypeError
 from taichi.lang.expr import Expr
@@ -20,7 +21,7 @@ def __init__(self, *args):
             args[i] = (args[i][0], ops.max(args[i][0], args[i][1]))
         for arg in args:
             for bound in arg:
-                if not isinstance(bound, int) and not (
+                if not isinstance(bound, (int, np.integer)) and not (
                         isinstance(bound, Expr)
                         and is_integral(bound.ptr.get_ret_type())):
                     raise TaichiTypeError(

diff --git a/tests/python/test_ndrange.py b/tests/python/test_ndrange.py
@@ -278,6 +278,18 @@ def example():
         example()
 
 
+@test_utils.test()
+def test_ndrange_should_accept_numpy_integer():
+    a, b = np.int64(0), np.int32(10)
+
+    @ti.kernel
+    def example():
+        for i in ti.ndrange((a, b)):
+            pass
+
+    example()
+
+
 @test_utils.test()
 def test_static_ndrange_non_integer_arguments():
     @ti.kernel