Save element_dim parameter for taichi.types.ndarray.

[mushroomfire]

Concisely describe the proposed feature
I would like to suggest saving element_dim parameter in taichi.types.ndarray data structure.

Describe the solution you'd like (if any)
Considering one want to interact with a 2-D numpy array and treat each row as a vector, and sometimes the length of row is variable. element_dim is suitable in such case while ndim and dtype maybe not a good choice. For example:
Generate a numpy array to inteprete the position in 2-D and 3-D.

import taichi as ti
import numpy as np
ti.init()
N = 10
pos_2 = np.random.random((N, 2)) 
pos_3 = np.random.random((N, 3))

One can handle pos_2 and pos_3 with the following kernel:

@ti.kernel
def vector_ele(arr:ti.types.ndarray(element_dim=1)) -> float:
    res = 0.
    for i in range(arr.shape[0]):
        res += (arr[i] - arr[0]).sum()
    return res

If only use ndim and dtype parameters, the code maybe like this:

@ti.kernel
def _vector_dtype(arr:ti.types.ndarray()) -> float:
    res = 0.
    for i in range(arr.shape[0]):
        res += (arr[i] - arr[0]).sum()
    return res

def vector_dtype(pos):
    dtype = ti.types.vector(n=pos.shape[1], dtype=float)
    pos_ti  = ti.ndarray(dtype=dtype, shape=(pos.shape[0]))
    pos_ti.from_numpy(pos)
    return _vector_dtype(pos_ti)

I have to transform numpy array to ti.ndarray first, and thus I think combination of ndim and dtype can not entirely replace the function of element_dim, which should be saved.

[bobcao3]

Or somehow make ti.template work with NDarrays (perhaps move element dim to templates as a type matching system, and if this can be compatible with fields this will be super extra helpful) because otherwise we would lose all the metaprogramming options with NDArrays.

[bobcao3]

@lin-hitonami @k-ye for visibility

[k-ye]

Or somehow make ti.template work with NDarrays (perhaps move element dim to templates as a type matching system, and if this can be compatible with fields this will be super extra helpful) because otherwise we would lose all the metaprogramming options with NDArrays.

I like this path. My nit on this is that instead of using ti.template(), let's stick with ti.types.ndarray() for clarity. When ti.types.ndarray() is used as a type annotation, it should support partial type matching. (https://docs.taichi-lang.org/docs/ndarray#using-taichi-ndarrays-in-tikernel) That is, if only ndim is specified, Taichi will try to infer dtype from the numpy array passed-in. WDYT?

I prefer not to add element_dim back, because it introduces things that is no longer orthogonal with the type system. What should we do when we specify ndarray(element_dim=2, dtype=ti.f32)?

[bobcao3]

Or somehow make ti.template work with NDarrays (perhaps move element dim to templates as a type matching system, and if this can be compatible with fields this will be super extra helpful) because otherwise we would lose all the metaprogramming options with NDArrays.

I like this path. My nit on this is that instead of using ti.template(), let's stick with ti.types.ndarray() for clarity. When ti.types.ndarray() is used as a type annotation, it should support partial type matching. (https://docs.taichi-lang.org/docs/ndarray#using-taichi-ndarrays-in-tikernel) That is, if only ndim is specified, Taichi will try to infer dtype from the numpy array passed-in. WDYT?

I prefer not to add element_dim back, because it introduces things that is no longer orthogonal with the type system. What should we do when we specify ndarray(element_dim=2, dtype=ti.f32)?

The use of template allow metaprogramming compatibility with fields. In terms of current CHI IR design this wouldn't work cleanly, but in terms of language design both NDarrays and Fields support indexed subscript interface and should be able to run the same templated code.

[k-ye]

Sure, ti.template() is more akin to C++'s T, whereas ti.types.ndarray() + partial matching is like an "Ndarray trait". I think "both NDarrays and Fields support indexed subscript interface" is also asking for traits.

[bobcao3]

Sure, ti.template() is more akin to C++'s T, whereas ti.types.ndarray() + partial matching is like an "Ndarray trait". I think "both NDarrays and Fields support indexed subscript interface" is also asking for traits.

Yes, I think we would need something along the lines of interface/traits

[ailzhang]

Hey @bobcao3 @mushroomfire @k-ye @strongoier , @jim19930609 and I had a quick discussion offline after the triage meeting, and we agree this is a pretty valid use case - no matter it's external array or not ;D

But the element_dim is kinda verbose and may confuse users, how about this? (let's keep ndarray case separate from ti.template for now until we finalize on what ti.template should do?)

@ti.kernel
def vector_ele(arr:ti.types.ndarray(dtype=ti.types.matrix())) -> float:
    res = 0.
    for i in range(arr.shape[0]):
        res += (arr[i] - arr[0]).sum()
    return res

[strongoier]

how about this?

Looks good. We can make all type annotations templates.

[k-ye]

+1 for the annotation syntax. I think it's a lot more clear to specify matrix, vector than an element_dim.

[mushroomfire]

Dose ti.types.matrix() work with a 3-D array? if we only want to treat the last dim as a vector. @ailzhang, @k-ye. element_dim can handle it.

Nf=100 # number of frames
N=1000 # number of particles
dim = 3 # dimension
trajec = np.random.random((Nf, N, dim))

BTW, I do think the element_dim is not clear, but if we want to drop out it, we should make sure the new parameters can entirely replace it with a more concise way.

[ailzhang]

Hey @mushroomfire currently ti.types.matrix represents 2d arrays specifically but we can add another sth like ti.types.tensor(ndim=3) to make it more clear? (name is not finalized, suggestions are welcome!

[mushroomfire]

Hey @ailzhang, I like ti.types.tensor(), but again, what should I do for a 4-D array, use ndim=4? can we ignore the ndim.

[ailzhang]

@mushroomfire yep you can ignore ndim, so sth like ti.types.tensor() can just accept array for any dimensions (as long as taichi kernel compiles) and ti.types.matrix would become an alias of ti.types.tensor(ndim=2).

[mushroomfire]

@ailzhang Yes, that is what I mean!

[strongoier]

@mushroomfire @ailzhang It seems that you are discussing dim with different meanings. Simply speaking, dtype=ti.types.matrix() is equivalent to the old element_dim=2 and dtype=ti.types.vector() is equivalent to the old element_dim=1.