Ndarray type annotations & error message refinement

[ailzhang]

We had a discussion about refining ndarray definition and type annotation in #4085 earlier this year and it's now a hard blocker if we want to properly release ndarray to users. Based on the discussions with @k-ye and @strongoier , we'll list out the desired APIs as well as some examples of bad error messages we wanna fix here.

Type annotation

Currently ti.types.ndarray() can take in dtype/element_dim/element_shape/field_dim as optional arguments, we'd like to:

• Rename field_dim to ndim (following numpy etc)
• Get rid of element_dim and element_shape in the type annotation

And the desired ndarray type annotation API may only takes two optional argument dtype and ndim, which looks like below:

arr: ti.types.ndarray()
arr: ti.types.ndarray(dtype=ti.f32)
arr: ti.types.ndarray(dtype=ti.math.vec3)
arr: ti.types.ndarray(ndim=2)
arr: ti.types.ndarray(dtype=ti.f32, ndim=2)

Error message

Currently we've got a lot to improve in terms of telling users what's wrong with their ndarray usage, listing a few examples below:

A running demo can be found at https://gist.github.com/ailzhang/d207a53c26720cccb7e7c5a9687009c2

tp_ivec3 = ti.types.vector(3, ti.i32)
tp_ivec2 = ti.types.vector(2, ti.i32)
tp_vec3 = ti.types.vector(3, ti.f32)

x = ti.ndarray(ti.f32, shape=(12, 13))
y = ti.ndarray(tp_ivec3, shape=(12,4))
z = ti.ndarray(ti.i32, shape=(12, 12))
m = ti.ndarray(tp_ivec2, shape=(2, 3))
n = ti.ndarray(ti.f32, shape=(12,))
k = ti.ndarray(tp_vec3, shape=(12,))

@ti.kernel
def test1(arr: ti.types.ndarray()):
    for I in ti.grouped(arr):
        arr[I] = 1

test1(x)
print(x.to_numpy())

# Case 1: The following currently errors out since the kernel implicitly assumes scalar element. We need better error message to users.
#test1(y)

@ti.kernel
def test2(arr: ti.types.ndarray(dtype=ti.f32)):
    for I in ti.grouped(arr):
        arr[I] = 2.0

test2(x)
print(x.to_numpy())

# Case 2: This works, but it should error out since it doesn't match type annotation. 
test2(z)
print(z.to_numpy())

# Case 3: this errors out, but currently error message should be mismatch dtype
# test2(y)

@ti.kernel
def test3(arr: ti.types.ndarray(dtype=tp_ivec3)):
    for I in ti.grouped(arr):
        arr[I] = [0, 1, 2]

test3(y)
print(y.to_numpy())

# Case 4: this errors out, but message should be mismatched dtype
# test3(m)

# Case 5: This should have errored out but didn't
# test3(k)

@ti.kernel
def test4(arr: ti.types.ndarray(dtype=tp_ivec3)):
    for I in ti.grouped(arr):
        arr[I] = [1, 2]

# Case 6: test4 is just bad kernel. It should just error out in compilation saying [1, 2] doesn't match ivec3 dtype
# test4(y)

@ti.kernel
def test5(arr: ti.types.ndarray(field_dim=1)):
    for i, j in arr:
        arr[i, j] = 0

# Case 7: test5 is just bad kernel, this should error out in compilation saying `for i, j in arr` doesn't match annotated field_dim=1 
# test5(x)

@ti.kernel
def test6(arr: ti.types.ndarray(field_dim=2)):
    for i, j in arr:
        arr[i, j] = 6

# Works as expected
test6(x)
print(x.to_numpy())

# Case 8: this errors out but message can be given ivec2 but arr requires scalar dtype. 
# test6(m)

# Case 9: this errors out but message can be give field_dim=1 but arr requires field_dim=2
# test6(n)

Note we'll need similar improvement when compiling kernels in AOT mode as well but those should be simple after the aforementioned issues are resolved.

Additional Cleanup

• We were not consistent in enforcing the order of arguments when defining a type or ndarray initialization, for example:

vec3 = ti.types.vector(3, ti.f32)
arr_ty = ti.types.ndarray(dtype=ti.f32, ndim=2)
arr = ti.ndarray(dtype=ti.f32, shape=(3, 3)) 

Notice the swap of dtype and shape arguments here? We'd actually like to unify to dtype first order for all of them, which involves adding some deprecation warnings.

• In addition, we should support vec/matrix dtype in ti.field() and deprecate ti.Vector.field() and ti.Matrix.field() APIs.

[turbo0628]

Yet another confusing error message

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

@ti.kernel
def add_num(num : ti.f32, arr : ti.types.ndarray(dtype=ti.math.mat3, ndim=2)):
    for i in ti.grouped(arr):
        arr[i] = arr[i] + num

arr_np = np.ones((2,2,3,3), dtype=np.float32)
arr_ti = ti.ndarray(dtype=ti.math.vec3, shape=(2, 2))
arr_ti.from_numpy(arr_np)
add_num(1.12, arr_ti)
print(arr_ti.to_numpy())

[Taichi] version 1.5.0, llvm 15.0.5, commit 6f8e0edb, osx, python 3.9.12
[Taichi] Starting on arch=metal
Traceback (most recent call last):
  File "/Users/lanhaidong/example_vec.py", line 13, in <module>
    arr_ti.from_numpy(arr_np)
  File "/Users/lanhaidong/work/taichi_fork/python/taichi/lang/util.py", line 310, in wrapped
    return func(*args, **kwargs)
  File "/Users/lanhaidong/work/taichi_fork/python/taichi/lang/matrix.py", line 1809, in from_numpy
    self._ndarray_matrix_from_numpy(arr, as_vector=1)
  File "/Users/lanhaidong/work/taichi_fork/python/taichi/lang/util.py", line 310, in wrapped
    return func(*args, **kwargs)
  File "/Users/lanhaidong/work/taichi_fork/python/taichi/lang/_ndarray.py", line 139, in _ndarray_matrix_from_numpy
    raise ValueError(
ValueError: Mismatch shape: (2, 2) expected, but (2, 2, 3, 3) provided

The "expected shape" hint doesn't include the element shape, so it's very confusing.