def _test_ndarray_numpy_io():
n = 7
m = 4
a = ti.ndarray(ti.i32, shape=(n, m))
a.fill(2)
b = ti.ndarray(ti.i32, shape=(n, m))
b.from_numpy(np.ones((n, m), dtype=np.int32) * 2)
assert (a.to_numpy() == b.to_numpy()).all()
d = 2
p = 4
x = ti.Vector.ndarray(d, ti.f32, p)
x.fill(2)
y = ti.Vector.ndarray(d, ti.f32, p)
y.from_numpy(np.ones((p, d), dtype=np.int32) * 2)
assert (x.to_numpy() == y.to_numpy()).all()
c = 2
d = 2
p = 4
x = ti.Matrix.ndarray(c, d, ti.f32, p)
x.fill(2)
y = ti.Matrix.ndarray(c, d, ti.f32, p)
y.from_numpy(np.ones((p, c, d), dtype=np.int32) * 2)
assert (x.to_numpy() == y.to_numpy()).all()
def test_ndarray_cuda_caching_allocator():
n = 8
a = ti.ndarray(ti.i32, shape=(n))
a.fill(2)
a = 1
b = ti.ndarray(ti.i32, shape=(n))
b.fill(2)
def test_ndarray():
x = ti.Vector.ndarray(dim, ti.f32, n_particles)
v = ti.Vector.ndarray(dim, ti.f32, n_particles)
C = ti.Matrix.ndarray(dim, dim, ti.f32, n_particles)
J = ti.ndarray(ti.f32, n_particles)
grid_v = ti.Vector.ndarray(dim, ti.f32, (n_grid, n_grid))
grid_m = ti.ndarray(ti.f32, (n_grid, n_grid))
run_test(x, v, C, J, grid_v, grid_m)
def test_ndarray_cuda_caching_allocator():
ti.init(arch=ti.cuda, ndarray_use_cached_allocator=True)
n = 8
a = ti.ndarray(ti.i32, shape=(n))
a.fill(2)
a = 1
b = ti.ndarray(ti.i32, shape=(n))
b.fill(2)
def _test_ndarray_numpy_io():
n = 7
m = 4
a = ti.ndarray(ti.i32, shape=(n, m))
a.fill(2)
b = ti.ndarray(ti.i32, shape=(n, m))
b.from_numpy(np.ones((n, m), dtype=np.int32) * 2)
assert (a.to_numpy() == b.to_numpy()).all()
def test_ndarray_as_template():
@ti.kernel
def func(arr_src: ti.template(), arr_dst: ti.template()):
for i, j in ti.ndrange(*arr_src.shape):
arr_dst[i, j] = arr_src[i, j]
arr_0 = ti.ndarray(ti.f32, shape=(5, 10))
arr_1 = ti.ndarray(ti.f32, shape=(5, 10))
with pytest.raises(ti.TaichiRuntimeTypeError,
match=r"Ndarray shouldn't be passed in via"):
func(arr_0, arr_1)
def test_opengl_exceed_max_ssbo():
# 8 ndarrays + args > 8 (maximum allowed)
n = 4
density1 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density2 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density3 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density4 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density5 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density6 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density7 = ti.ndarray(dtype=ti.f32, shape=(n, n))
density8 = ti.ndarray(dtype=ti.f32, shape=(n, n))
@ti.kernel
def init(d: ti.i32, density1: ti.types.ndarray(),
density2: ti.types.ndarray(), density3: ti.types.ndarray(),
density4: ti.types.ndarray(), density5: ti.types.ndarray(),
density6: ti.types.ndarray(), density7: ti.types.ndarray(),
density8: ti.types.ndarray()):
for i, j in density1:
density1[i, j] = d + 1
density2[i, j] = d + 2
density3[i, j] = d + 3
density4[i, j] = d + 4
density5[i, j] = d + 5
density6[i, j] = d + 6
density7[i, j] = d + 7
density8[i, j] = d + 8
with pytest.raises(RuntimeError):
init(0, density1, density2, density3, density4, density5, density6,
density7, density8)
def test_different_shape():
n1 = 4
x = ti.ndarray(dtype=ti.f32, shape=(n1, n1))
@ti.kernel
def init(d: ti.i32, arr: ti.types.ndarray()):
for i, j in arr:
arr[i, j] = d
init(2, x)
assert (x.to_numpy() == (np.ones(shape=(n1, n1)) * 2)).all()
n2 = 8
y = ti.ndarray(dtype=ti.f32, shape=(n2, n2))
init(3, y)
assert (y.to_numpy() == (np.ones(shape=(n2, n2)) * 3)).all()
def test_matrix_int():
n = 4
A = ti.Matrix([4, 5] * n)
res = ti.ndarray(ti.i32, shape=(1, ))
graph = build_graph_matrix(n, dtype=ti.i32)
graph.run({"mat": A, "res": res})
assert (res.to_numpy()[0] == 36)
def test_matrix_float():
n = 4
A = ti.Matrix([4.2, 5.7] * n)
res = ti.ndarray(ti.f32, shape=(1))
graph = build_graph_matrix(n, dtype=ti.f32)
graph.run({"mat": A, "res": res})
assert res.to_numpy()[0] == test_utils.approx(39.6, rel=1e-5)
def test_vector_int():
n = 12
A = ti.Vector([1, 3, 13, 4, 5, 6, 7, 2, 3, 4, 1, 25])
res = ti.ndarray(ti.i32, shape=(1, ))
graph = build_graph_vector(n, dtype=ti.i32)
graph.run({"mat": A, "res": res})
assert (res.to_numpy()[0] == 87)
def test_vector_float():
n = 8
A = ti.Vector([1.4, 3.7, 13.2, 4.5, 5.6, 6.1, 7.2, 2.6])
res = ti.ndarray(ti.f32, shape=(1, ))
graph = build_graph_vector(n, dtype=ti.f32)
graph.run({"mat": A, "res": res})
assert res.to_numpy()[0] == test_utils.approx(57.5, rel=1e-5)
def _test_size_in_bytes():
a = ti.ndarray(ti.i32, 8)
assert a._get_element_size() == 4
assert a._get_nelement() == 8
b = ti.Vector.ndarray(10, ti.f64, 5)
assert b._get_element_size() == 8
assert b._get_nelement() == 50
def test_subscript():
a = ti.ndarray(ti.i32, shape=(10, 10))
@ti.kernel
def any_array(x: ti.any_arr()):
b = x[3, 1.1]
with pytest.raises(TypeError, match="indices must be integers"):
any_array(a)
def test_ndarray_compound_element():
n = 10
a = ti.ndarray(ti.i32, shape=(n, ))
vec3 = ti.types.vector(3, ti.i32)
b = ti.ndarray(vec3, shape=(n, n))
assert isinstance(b, ti.MatrixNdarray)
assert b.shape == (n, n)
assert b.element_type.dtype == ti.i32
assert b.element_type.shape == (3, 1)
matrix34 = ti.types.matrix(3, 4, float)
c = ti.ndarray(matrix34, shape=(n, n + 1), layout=ti.Layout.SOA)
assert isinstance(c, ti.MatrixNdarray)
assert c.shape == (n, n + 1)
assert c.element_type.dtype == ti.f32
assert c.element_type.shape == (3, 4)
assert c.layout == ti.Layout.SOA
def _test_scalar_ndarray(dtype, shape):
x = ti.ndarray(dtype, shape)
if isinstance(shape, tuple):
assert x.shape == shape
else:
assert x.shape == (shape, )
assert x.element_shape == ()
assert x.dtype == dtype
def test_get_external_tensor_shape_access_ndarray(size):
@ti.kernel
def func(x: ti.any_arr(), index: ti.template()) -> ti.i32:
return x.shape[index]
x_hat = ti.ndarray(ti.i32, shape=size)
for idx, y_ref in enumerate(size):
y_hat = func(x_hat, idx)
assert y_ref == y_hat, "Size of axis {} should equal {} and not {}.".format(
idx, y_ref, y_hat)
def test_paddle_view():
@ti.kernel
def copy(x: ti.types.ndarray(), y: ti.types.ndarray()):
for i, j in x:
y[i, j] = x[i, j]
x = paddle.to_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).T
y = ti.ndarray(int, (3, 3))
copy(x, y)
def test_torch_view():
@ti.kernel
def copy(x: ti.types.ndarray(), y: ti.types.ndarray()):
for i, j in x:
y[i, j] = x[i, j]
x = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).T
y = ti.ndarray(int, (3, 3))
with pytest.raises(ValueError,
match=r'Non contiguous tensors are not supported'):
copy(x, y)
def test_opengl_8_ssbo():
# 6 ndarrays + gtmp + args
n = 4
density1 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
density2 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
density3 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
density4 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
density5 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
density6 = ti.ndarray(dtype=ti.f32, shape=(4, 4))
@ti.kernel
def init(d: ti.i32, density1: ti.types.ndarray(),
density2: ti.types.ndarray(), density3: ti.types.ndarray(),
density4: ti.types.ndarray(), density5: ti.types.ndarray(),
density6: ti.types.ndarray()):
for i, j in density1:
density1[i, j] = d + 1
density2[i, j] = d + 2
density3[i, j] = d + 3
density4[i, j] = d + 4
density5[i, j] = d + 5
density6[i, j] = d + 6
init(0, density1, density2, density3, density4, density5, density6)
assert (density1.to_numpy() == (np.zeros(shape=(n, n)) + 1)).all()
assert (density2.to_numpy() == (np.zeros(shape=(n, n)) + 2)).all()
assert (density3.to_numpy() == (np.zeros(shape=(n, n)) + 3)).all()
assert (density4.to_numpy() == (np.zeros(shape=(n, n)) + 4)).all()
assert (density5.to_numpy() == (np.zeros(shape=(n, n)) + 5)).all()
assert (density6.to_numpy() == (np.zeros(shape=(n, n)) + 6)).all()
def _test_ndarray_deepcopy():
n = 16
x = ti.ndarray(ti.i32, shape=n)
x[0] = 1
x[4] = 2
y = copy.deepcopy(x)
assert y.shape == x.shape
assert y.dtype == x.dtype
assert y[0] == 1
assert y[4] == 2
x[0] = 4
x[4] = 5
assert y[0] == 1
assert y[4] == 2
x = ti.Vector.ndarray(10, ti.i32, 5, layout=ti.Layout.SOA)
x[1][0] = 4
x[2][4] = 5
y = copy.deepcopy(x)
assert y.shape == x.shape
assert y.dtype == x.dtype
assert y.n == x.n
assert y.layout == x.layout
assert y[1][0] == 4
assert y[2][4] == 5
x[1][0] = 1
x[2][4] = 2
assert y[1][0] == 4
assert y[2][4] == 5
x = ti.Matrix.ndarray(2, 2, ti.i32, 5, layout=ti.Layout.AOS)
x[0][0, 0] = 7
x[4][1, 0] = 9
y = copy.deepcopy(x)
assert y.shape == x.shape
assert y.dtype == x.dtype
assert y.m == x.m
assert y.n == x.n
assert y.layout == x.layout
assert y[0][0, 0] == 7
assert y[4][1, 0] == 9
x[0][0, 0] = 3
x[4][1, 0] = 5
assert y[0][0, 0] == 7
assert y[4][1, 0] == 9
def test_torch_based_ndarray_not_supported():
x = ti.ndarray(ti.f32, shape=(4, 2))
@ti.kernel
def init(x: ti.any_arr()):
for i, j in x:
x[i, j] = i + j
with pytest.raises(
AssertionError,
match=
r'Torch-based ndarray is only supported on taichi x64/arm64/cuda backend'
):
init(x)
def _test_ndarray_copy_from_ndarray():
n = 16
a = ti.ndarray(ti.i32, shape=n)
b = ti.ndarray(ti.i32, shape=n)
a[0] = 1
a[4] = 2
b[0] = 4
b[4] = 5
a.copy_from(b)
assert a[0] == 4
assert a[4] == 5
x = ti.Vector.ndarray(10, ti.i32, 5, layout=ti.Layout.SOA)
y = ti.Vector.ndarray(10, ti.i32, 5, layout=ti.Layout.SOA)
x[1][0] = 1
x[2][4] = 2
y[1][0] = 4
y[2][4] = 5
x.copy_from(y)
assert x[1][0] == 4
assert x[2][4] == 5
x = ti.Matrix.ndarray(2, 2, ti.i32, 5, layout=ti.Layout.AOS)
y = ti.Matrix.ndarray(2, 2, ti.i32, 5, layout=ti.Layout.AOS)
x[0][0, 0] = 1
x[4][1, 0] = 3
y[0][0, 0] = 4
y[4][1, 0] = 6
x.copy_from(y)
assert x[0][0, 0] == 4
assert x[4][1, 0] == 6
def test_aot_ndarray_template_mixed():
@ti.kernel
def run(arr: ti.types.ndarray(), val1: ti.f32, val2: ti.template()):
for i in arr:
arr[i] = val1 + val2
with tempfile.TemporaryDirectory() as tmpdir:
x = ti.ndarray(dtype=ti.f32, shape=16)
m = ti.aot.Module(ti.opengl)
m.add_kernel(run, template_args={'arr': x, 'val2': 42})
m.save(tmpdir, '')
with open(os.path.join(tmpdir, 'metadata.json')) as json_file:
res = json.load(json_file)
args_count = res['aot_data']['kernels']['run']['args_count']
assert args_count == 2, res # `arr` and `val1`
def test_arg_mismatched_field_dim_ndarray():
n = 4
@ti.kernel
def test(pos: ti.types.ndarray(field_dim=1)):
for i in range(n):
pos[i] = 2.5
sym_pos = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'pos', ti.f32, 1)
g_init = ti.graph.GraphBuilder()
g_init.dispatch(test, sym_pos)
g = g_init.compile()
a = ti.ndarray(ti.f32, shape=(n, n))
with pytest.raises(RuntimeError, match="Dispatch node is compiled for"):
g.run({'pos': a})
def test_ndarray_int():
n = 4
@ti.kernel
def test(pos: ti.types.ndarray(field_dim=1)):
for i in range(n):
pos[i] = 1
sym_pos = ti.graph.Arg(ti.graph.ArgKind.NDARRAY, 'pos', ti.i32)
g_init = ti.graph.GraphBuilder()
g_init.dispatch(test, sym_pos)
g = g_init.compile()
a = ti.ndarray(ti.i32, shape=(n, ))
g.run({'pos': a})
assert (a.to_numpy() == np.ones(4)).all()
def test_ndarray_0dim():
@ti.kernel
def test(pos: ti.types.ndarray(dtype=ti.i32, field_dim=0)):
pos[None] = 1
sym_pos = ti.graph.Arg(ti.graph.ArgKind.NDARRAY,
'pos',
ti.i32,
field_dim=0)
g_init = ti.graph.GraphBuilder()
g_init.dispatch(test, sym_pos)
g = g_init.compile()
a = ti.ndarray(ti.i32, shape=())
g.run({'pos': a})
assert a.to_numpy() == 1
def compile_kernel_aot(arch):
ti.init(arch=arch)
@ti.kernel
def run(base: int, arr: ti.types.ndarray()):
for i in arr:
arr[i] = base + i
arr = ti.ndarray(int, shape=16)
assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
m = ti.aot.Module(arch)
m.add_kernel(run, template_args={'arr': arr})
m.save(dir_name, 'whatever')
def test_aot_ndarray_range_hint():
density = ti.ndarray(dtype=ti.f32, shape=(8, 8))
@ti.kernel
def init(density: ti.types.ndarray()):
for i, j in density:
density[i, j] = 1
with tempfile.TemporaryDirectory() as tmpdir:
m = ti.aot.Module(ti.opengl)
m.add_kernel(init, template_args={'density': density})
m.save(tmpdir, '')
with open(os.path.join(tmpdir, 'metadata.json')) as json_file:
res = json.load(json_file)
range_hint = res['aot_data']['kernels']['init']['tasks'][0][
'range_hint']
assert range_hint == 'arg 0'
def test_ndarray_1d():
n = 4
@ti.kernel
def run(x: ti.types.ndarray(), y: ti.types.ndarray()):
for i in range(n):
x[i] += i + y[i]
a = ti.ndarray(ti.i32, shape=(n, ))
for i in range(n):
a[i] = i * i
b = np.ones((n, ), dtype=np.int32)
run(a, b)
for i in range(n):
assert a[i] == i * i + i + 1
run(b, a)
for i in range(n):
assert b[i] == i * i + (i + 1) * 2