[Bug] [lang] Allow numpy int as snode dimension

6186.py

@@ -0,0 +1,7 @@
+import taichi as ti
+import numpy as np
+
+ti.init(arch=ti.gpu)
+
+np_int = np.max([0, 1, 2])
+block = ti.root.pointer(ti.i, np_int)

python/taichi/lang/snode.py

@@ -30,7 +30,7 @@ def dense(self, axes, dimensions):
         Returns:
             The added :class:`~taichi.lang.SNode` instance.
         """
-        if isinstance(dimensions, int):
+        if isinstance(dimensions, numbers.Number):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.dense(axes, dimensions,
@@ -46,7 +46,7 @@ def pointer(self, axes, dimensions):
         Returns:
             The added :class:`~taichi.lang.SNode` instance.
         """
-        if isinstance(dimensions, int):
+        if isinstance(dimensions, numbers.Number):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.pointer(axes, dimensions,
@@ -90,7 +90,7 @@ def bitmasked(self, axes, dimensions):
         Returns:
             The added :class:`~taichi.lang.SNode` instance.
         """
-        if isinstance(dimensions, int):
+        if isinstance(dimensions, numbers.Number):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.bitmasked(axes, dimensions,
@@ -107,7 +107,7 @@ def quant_array(self, axes, dimensions, max_num_bits):
         Returns:
             The added :class:`~taichi.lang.SNode` instance.
         """
-        if isinstance(dimensions, int):
+        if isinstance(dimensions, numbers.Number):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.quant_array(axes, dimensions, max_num_bits,

tests/python/test_fields_builder.py

@@ -1,6 +1,7 @@
 import pytest
 from taichi.lang.exception import TaichiRuntimeError
 
+import numpy as np
 import taichi as ti
 from tests import test_utils
 
@@ -227,3 +228,14 @@ def calc_loss(arr: ti.template(), loss: ti.template()):
     mul.grad(arr, out)
     for i in range(10):
         assert arr.grad[i] == 2.0
+
+
+@test_utils.test(arch=ti.cpu)
+def test_fields_builder_numpy_dimension():
+    shape = np.int32(5)
+    fb = ti.FieldsBuilder()
+    x = ti.field(ti.f32)
+    y = ti.field(ti.i32)
+    fb.dense(ti.i, shape).place(x)
+    fb.pointer(ti.j, shape).place(y)
+    fb.finalize()