def grad_test(tifunc, npfunc=None):
npfunc = npfunc or tifunc
print(
f'arch={ti.lang.impl.current_cfg().arch} default_fp={ti.lang.impl.current_cfg().default_fp}'
)
x = ti.field(ti.lang.impl.current_cfg().default_fp)
y = ti.field(ti.lang.impl.current_cfg().default_fp)
ti.root.dense(ti.i, 1).place(x, x.grad, y, y.grad)
@ti.kernel
def func():
for i in x:
y[i] = tifunc(x[i])
v = 0.234
y.grad[0] = 1
x[0] = v
func()
func.grad()
assert y[0] == test_utils.approx(npfunc(v), rel=1e-4)
assert x.grad[0] == test_utils.approx(grad(npfunc)(v), rel=1e-4)
def test_ad_reduce_fwd():
N = 16
x = ti.field(dtype=ti.f32, shape=N)
loss = ti.field(dtype=ti.f32, shape=())
ti.root.lazy_dual()
@ti.kernel
def func():
for i in x:
loss[None] += x[i]**2
total_loss = 0
for i in range(N):
x[i] = i
total_loss += i * i
with ti.ad.FwdMode(loss=loss, parameters=x, seed=[1.0 for _ in range(N)]):
func()
assert total_loss == test_utils.approx(loss[None])
sum = 0
for i in range(N):
sum += i * 2
assert loss.dual[None] == test_utils.approx(sum)
def _test_polar_decomp(dim, dt):
m = ti.Matrix.field(dim, dim, dt)
r = ti.Matrix.field(dim, dim, dt)
s = ti.Matrix.field(dim, dim, dt)
I = ti.Matrix.field(dim, dim, dt)
D = ti.Matrix.field(dim, dim, dt)
ti.root.place(m, r, s, I, D)
@ti.kernel
def polar():
R, S = ti.polar_decompose(m[None], dt)
r[None] = R
s[None] = S
m[None] = R @ S
I[None] = R @ R.transpose()
D[None] = S - S.transpose()
def V(i, j):
return i * 2 + j * 7 + int(i == j) * 3
for i in range(dim):
for j in range(dim):
m[None][i, j] = V(i, j)
polar()
tol = 5e-5 if dt == ti.f32 else 1e-12
for i in range(dim):
for j in range(dim):
assert m[None][i, j] == test_utils.approx(V(i, j), abs=tol)
assert I[None][i, j] == test_utils.approx(int(i == j), abs=tol)
assert D[None][i, j] == test_utils.approx(0, abs=tol)
def test_matrix_factories():
a = ti.Vector.field(3, dtype=ti.i32, shape=3)
b = ti.Matrix.field(2, 2, dtype=ti.f32, shape=2)
c = ti.Matrix.field(2, 3, dtype=ti.f32, shape=2)
@ti.kernel
def fill():
b[0] = ti.Matrix.identity(ti.f32, 2)
b[1] = ti.Matrix.rotation2d(math.pi / 3)
c[0] = ti.Matrix.zero(ti.f32, 2, 3)
c[1] = ti.Matrix.one(ti.f32, 2, 3)
for i in ti.static(range(3)):
a[i] = ti.Vector.unit(3, i)
fill()
for i in range(3):
for j in range(3):
assert a[i][j] == int(i == j)
sqrt3o2 = math.sqrt(3) / 2
assert b[0].to_numpy() == test_utils.approx(np.eye(2))
assert b[1].to_numpy() == test_utils.approx(
np.array([[0.5, -sqrt3o2], [sqrt3o2, 0.5]]))
assert c[0].to_numpy() == test_utils.approx(np.zeros((2, 3)))
assert c[1].to_numpy() == test_utils.approx(np.ones((2, 3)))
def _test_closing_offline_cache_for_a_kernel(curr_arch, kernel, args, result):
count_of_cache_file = len(listdir(tmp_offline_cache_file_path()))
ti.init(arch=curr_arch,
enable_fallback=False,
offline_cache=False,
offline_cache_file_path=tmp_offline_cache_file_path())
res1 = kernel(*args)
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == get_expected_num_cache_files()
ti.init(arch=curr_arch,
enable_fallback=False,
offline_cache=False,
offline_cache_file_path=tmp_offline_cache_file_path())
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == get_expected_num_cache_files()
res2 = kernel(*args)
assert res1 == test_utils.approx(result) and res1 == test_utils.approx(
res2)
ti.reset()
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == get_expected_num_cache_files()
def helper():
x = ti.field(dtype=ti.f32, shape=())
y = ti.field(dtype=ti.f32, shape=())
x[None] = 3.14
y[None] = 4.14
assert x[None] == test_utils.approx(3.14)
assert y[None] == test_utils.approx(4.14)
x[None] = 6.28
y[None] = 7.28
assert x[None] == test_utils.approx(6.28)
assert y[None] == test_utils.approx(7.28)
def test_precision():
u = ti.field(ti.f64, shape=())
v = ti.field(ti.f64, shape=())
w = ti.field(ti.f64, shape=())
@ti.kernel
def forward():
v[None] = ti.sqrt(ti.cast(u[None] + 3.25, ti.f64))
w[None] = ti.cast(u[None] + 7, ti.f64) / ti.cast(u[None] + 3, ti.f64)
forward()
assert v[None]**2 == test_utils.approx(3.25, abs=1e-12)
assert w[None] * 3 == test_utils.approx(7, abs=1e-12)
def test_unary_op():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo():
x[None] = -y[None]
x[None] = ti.floor(x[None])
y[None] = ti.ceil(y[None])
y[None] = -1.4
foo()
assert (x[None] == test_utils.approx(1, rel=1e-3))
assert (y[None] == test_utils.approx(-1, rel=1e-3))
def test_expr_dict_basic():
@ti.kernel
def func(u: int, v: float) -> float:
x = {'foo': 2 + u, 'bar': 3 + v}
return x['foo'] * 100 + x['bar']
assert func(2, 0.1) == test_utils.approx(403.1)
def _test_svd(dt, n):
print(
f'arch={ti.lang.impl.current_cfg().arch} default_fp={ti.lang.impl.current_cfg().default_fp} fast_math={ti.lang.impl.current_cfg().fast_math} dim={n}'
)
A = ti.Matrix.field(n, n, dtype=dt, shape=())
A_reconstructed = ti.Matrix.field(n, n, dtype=dt, shape=())
U = ti.Matrix.field(n, n, dtype=dt, shape=())
UtU = ti.Matrix.field(n, n, dtype=dt, shape=())
sigma = ti.Matrix.field(n, n, dtype=dt, shape=())
V = ti.Matrix.field(n, n, dtype=dt, shape=())
VtV = ti.Matrix.field(n, n, dtype=dt, shape=())
@ti.kernel
def run():
U[None], sigma[None], V[None] = ti.svd(A[None], dt)
UtU[None] = U[None].transpose() @ U[None]
VtV[None] = V[None].transpose() @ V[None]
A_reconstructed[None] = U[None] @ sigma[None] @ V[None].transpose()
if n == 3:
A[None] = [[1, 1, 3], [9, -3, 2], [-3, 4, 2]]
else:
A[None] = [[1, 1], [2, 3]]
run()
tol = 1e-5 if dt == ti.f32 else 1e-12
assert mat_equal(UtU.to_numpy(), np.eye(n), tol=tol)
assert mat_equal(VtV.to_numpy(), np.eye(n), tol=tol)
assert mat_equal(A_reconstructed.to_numpy(), A.to_numpy(), tol=tol)
for i in range(n):
for j in range(n):
if i != j:
assert sigma[None][i, j] == test_utils.approx(0)
def test_ad_frac():
@ti.func
def frac(x):
fractional = x - ti.floor(x) if x > 0. else x - ti.ceil(x)
return fractional
@ti.kernel
def ti_frac(input_field: ti.template(), output_field: ti.template()):
for i in input_field:
output_field[i] = frac(input_field[i])**2
@ti.kernel
def calc_loss(input_field: ti.template(), loss: ti.template()):
for i in input_field:
loss[None] += input_field[i]
n = 10
field0 = ti.field(dtype=ti.f32, shape=(n, ), needs_grad=True)
randoms = np.random.randn(10).astype(np.float32)
field0.from_numpy(randoms)
field1 = ti.field(dtype=ti.f32, shape=(n, ), needs_grad=True)
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
with ti.Tape(loss):
ti_frac(field0, field1)
calc_loss(field1, loss)
grads = field0.grad.to_numpy()
expected = np.modf(randoms)[0] * 2
for i in range(n):
assert grads[i] == test_utils.approx(expected[i], rel=1e-4)
def test_expr_list_basic():
@ti.kernel
def func(u: int, v: float) -> float:
x = [2 + u, 3 + v]
return x[0] * 100 + x[1]
assert func(1, 1.1) == test_utils.approx(304.1)
def test_scalar_argument():
@ti.kernel
def add(a: ti.f32, b: ti.f32) -> ti.f32:
a = a + b
return a
assert add(1.0, 2.0) == test_utils.approx(3.0)
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_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_constant_matrices():
assert ti.cos(math.pi / 3) == test_utils.approx(0.5)
assert np.allclose((-ti.Vector([2, 3])).to_numpy(), np.array([-2, -3]))
assert ti.cos(ti.Vector([2, 3])).to_numpy() == test_utils.approx(
np.cos(np.array([2, 3])))
assert ti.max(2, 3) == 3
res = ti.max(4, ti.Vector([3, 4, 5]))
assert np.allclose(res.to_numpy(), np.array([4, 4, 5]))
res = ti.Vector([2, 3]) + ti.Vector([3, 4])
assert np.allclose(res.to_numpy(), np.array([5, 7]))
res = ti.atan2(ti.Vector([2, 3]), ti.Vector([3, 4]))
assert res.to_numpy() == test_utils.approx(
np.arctan2(np.array([2, 3]), np.array([3, 4])))
res = ti.Matrix([[2, 3], [4, 5]]) @ ti.Vector([2, 3])
assert np.allclose(res.to_numpy(), np.array([13, 23]))
v = ti.Vector([3, 4])
w = ti.Vector([5, -12])
r = ti.Vector([1, 2, 3, 4])
s = ti.Matrix([[1, 2], [3, 4]])
assert v.normalized().to_numpy() == test_utils.approx(np.array([0.6, 0.8]))
assert v.cross(w) == test_utils.approx(-12 * 3 - 4 * 5)
w.y = v.x * w[0]
r.x = r.y
r.y = r.z
r.z = r.w
r.w = r.x
assert np.allclose(w.to_numpy(), np.array([5, 15]))
assert ti.select(ti.Vector([1, 0]), ti.Vector([2, 3]),
ti.Vector([4, 5])) == ti.Vector([2, 5])
s[0, 1] = 2
assert s[0, 1] == 2
@ti.kernel
def func(t: ti.i32):
m = ti.Matrix([[2, 3], [4, t]])
print(m @ ti.Vector([2, 3]))
m += ti.Matrix([[3, 4], [5, t]])
print(m @ v)
print(r.x, r.y, r.z, r.w)
s = w.transpose() @ m
print(s)
print(m)
func(5)
def test_grouped_static_for_cast():
@ti.kernel
def foo() -> ti.f32:
ret = 0.
for I in ti.static(ti.grouped(ti.ndrange((4, 5), (3, 5), 5))):
tmp = I.cast(float)
ret += tmp[2] / 2
return ret
assert foo() == test_utils.approx(10)
def test_image_resize_sum(resx, resy, comp, scale):
shape = (resx, resy)
if comp != 1:
shape = shape + (comp, )
old_img = np.random.rand(*shape).astype(np.float32)
if resx == resy:
new_img = ti.imresize(old_img, resx * scale)
else:
new_img = ti.imresize(old_img, resx * scale, resy * scale)
assert np.sum(old_img) * scale**2 == test_utils.approx(np.sum(new_img))
def test_random_vector_dup_eval():
a = ti.Vector.field(2, ti.f32, ())
@ti.kernel
def func():
a[None] = ti.Vector([ti.random(), 1]).normalized()
for i in range(4):
func()
assert a[None].norm_sqr() == test_utils.approx(1)
def test_expr_dict_field():
a = ti.field(ti.f32, shape=(4, ))
@ti.kernel
def func() -> float:
x = {'foo': 2 + a[0], 'bar': 3 + a[1]}
return x['foo'] * 100 + x['bar']
a[0] = 2
a[1] = 0.1
assert func() == test_utils.approx(403.1)
def test_func_random_argument_dup_eval():
@ti.func
def func(a):
return ti.Vector([ti.cos(a), ti.sin(a)])
@ti.kernel
def kern() -> ti.f32:
return func(ti.random()).norm_sqr()
for i in range(4):
assert kern() == test_utils.approx(1.0, rel=5e-5)
def test_arg_f16():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo(a: ti.f16):
x[None] = y[None] + a
y[None] = -0.3
foo(1.2)
assert (x[None] == test_utils.approx(0.9, rel=1e-3))
def test_extra_unary_promote():
dtype = ti.f16
x = ti.field(dtype, shape=())
y = ti.field(dtype, shape=())
@ti.kernel
def foo():
x[None] = abs(y[None])
y[None] = -0.3
foo()
assert (x[None] == test_utils.approx(0.3, rel=1e-3))
def test_binary_extra_promote():
x = ti.field(dtype=ti.f16, shape=())
y = ti.field(dtype=ti.f16, shape=())
z = ti.field(dtype=ti.f16, shape=())
@ti.kernel
def foo():
y[None] = x[None]**2
z[None] = ti.atan2(y[None], 0.3)
x[None] = 0.1
foo()
assert (z[None] == test_utils.approx(math.atan2(0.1**2, 0.3), rel=1e-3))
def _test_offline_cache_for_a_kernel(curr_arch, kernel, args, result):
count_of_cache_file = len(listdir(tmp_offline_cache_file_path()))
ti.init(arch=curr_arch,
enable_fallback=False,
**current_thread_ext_options())
res1 = kernel(*args)
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == 0 * cache_files_num_per_kernel
ti.init(arch=curr_arch,
enable_fallback=False,
**current_thread_ext_options())
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == 1 * cache_files_num_per_kernel
res2 = kernel(*args)
assert res1 == test_utils.approx(result) and res1 == test_utils.approx(
res2)
ti.reset()
assert len(listdir(tmp_offline_cache_file_path())
) - count_of_cache_file == 1 * cache_files_num_per_kernel
def test_const_func_ret():
ti.init()
@ti.kernel
def func1() -> ti.f32:
return 3
@ti.kernel
def func2() -> ti.i32:
return 3.3 # return type mismatch, will be auto-casted into ti.i32
assert func1() == test_utils.approx(3)
assert func2() == 3
def test_ad_reduce():
N = 16
x = ti.field(dtype=ti.f32, shape=N, needs_grad=True)
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
@ti.kernel
def func():
for i in x:
loss[None] += x[i]**2
total_loss = 0
for i in range(N):
x[i] = i
total_loss += i * i
loss.grad[None] = 1
func()
func.grad()
assert total_loss == test_utils.approx(loss[None])
for i in range(N):
assert x.grad[i] == test_utils.approx(i * 2)
def test_minimization():
from taichi.examples.autodiff.minimization import (L, gradient_descent, n,
reduce, x, y)
for i in range(n):
x[i] = random.random()
y[i] = random.random()
for k in range(100):
with ti.Tape(loss=L):
reduce()
gradient_descent()
for i in range(n):
assert x[i] == test_utils.approx(y[i], rel=1e-2)
def test_diag():
m1 = ti.Matrix.field(3, 3, dtype=ti.f32, shape=())
@ti.kernel
def fill():
m1[None] = ti.Matrix.diag(dim=3, val=1.4)
fill()
for i in range(3):
for j in range(3):
if i == j:
assert m1[None][i, j] == test_utils.approx(1.4)
else:
assert m1[None][i, j] == 0.0
def test_random_independent_product():
n = 1024
x = ti.field(ti.f32, shape=n * n)
@ti.kernel
def fill():
for i in range(n * n):
a = ti.random()
b = ti.random()
x[i] = a * b
fill()
X = x.to_numpy()
for i in range(4):
assert X.mean() == test_utils.approx(1 / 4, rel=1e-2)