Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
@ti.all_archs
def test_linear_nested_aos():
x = ti.var(ti.i32)
y = ti.var(ti.i32)
n = 128
@ti.layout
def place():
ti.root.dense(ti.i, n // 16).dense(ti.i, 16).place(x, y)
for i in range(n):
x[i] = i
y[i] = i + 123
for i in range(n):
assert x[i] == i
@ti.all_archs
def test_f64():
val = ti.var(ti.f64)
n = 4
m = 7
@ti.layout
def values():
ti.root.dense(ti.ij, (n, m)).place(val)
for i in range(n):
for j in range(m):
val[i, j] = (i + j * 3) * 1e100
val.from_numpy(val.to_numpy() * 2)
@ti.all_archs
def test_abs():
@ti.kernel
def func():
print(1)
func()
@ti.all_archs
def test_loop_grad_complex():
return # This case is not supported yet
x = ti.var(ti.f32)
n = 16
m = 8
@ti.layout
def place():
ti.root.dense(ti.ij, (n, m)).place(x)
ti.root.lazy_grad()
@ti.kernel
def func():
for k in range(n):
t = k * k
@ti.all_archs
def test_pointer():
if ti.get_os_name() == 'win':
# This test not supported on Windows due to the VirtualAlloc issue #251
return
x = ti.var(ti.f32)
s = ti.var(ti.i32)
n = 128
@ti.layout
def place():
ti.root.dense(ti.i, n).pointer().dense(ti.i, n).place(x)
ti.root.place(s)
@ti.kernel
def activate():
@ti.all_archs
def test_io():
ti.cfg.arch = ti.cuda
x = ti.var(ti.i32)
n = 128
@ti.layout
def place():
ti.root.dense(ti.i, n).place(x)
x[3] = 123
x[4] = 456
assert x[3] == 123
assert x[4] == 456
@ti.all_archs
def test_numpy_2d_transpose():
val = ti.var(ti.i32)
n = 8
m = 8
@ti.layout
def values():
ti.root.dense(ti.ij, (n, m)).place(val)
@ti.kernel
def test_numpy(arr: ti.ext_arr()):
for i in ti.grouped(val):
val[i] = arr[i]
a = np.empty(shape=(n, m), dtype=np.int32)
@ti.all_archs
def test_ad_reduce():
x = ti.var(ti.f32)
loss = ti.var(ti.f32)
N = 16
@ti.layout
def place():
ti.root.place(loss, loss.grad).dense(ti.i, N).place(x, x.grad)
@ti.kernel
def func():
for i in x:
loss.atomic_add(ti.sqr(x[i]))
total_loss = 0
@ti.all_archs
def test_simple():
x = ti.var(ti.i32)
n = 128
@ti.layout
def place():
ti.root.dense(ti.i, n).place(x)
@ti.kernel
def func():
x[7] = 120
func()
for i in range(n):
@ti.all_archs
def test_kernel_template_mapper_numpy():
x = ti.var(ti.i32)
y = ti.var(ti.f32)
@ti.layout
def layout():
ti.root.place(x, y)
annotations = (ti.template(), ti.template(), ti.ext_arr())
import numpy as np
mapper = ti.KernelTemplateMapper(annotations, (0, 1, 2))
assert mapper.lookup((0, 0, np.ones(shape=(1, 2, 3), dtype=np.float32))) == 0
assert mapper.lookup((0, 0, np.ones(shape=(1, 2, 4), dtype=np.float32))) == 0
assert mapper.lookup((0, 0, np.ones(shape=(1, 2, 1), dtype=np.int32))) == 1