taichi怎么加速np.where和vstack这种矩阵操作呢?

前一段时间兴致冲冲试了一下TAICHI,结果一直报错,后来发现我的运算里面有矩阵切片操作,而太极没法用矩阵切片,请问大佬,对于矩阵切片和类似vstack没有办法实现吗

Hi 0crazy0ac0,
Taichi支持Matrix类型的切片,但是不支持Ndarray类型的切片。Matrix的切片比如:

@ti.kernel
def assign_col() -> ti.types.matrix(3, 4, ti.i32):
    mat = ti.Matrix([[0, 0, 0, 0] for _ in range(3)])
    col = ti.Vector([1, 2, 3]) 
    mat[:, 0] = col 
    return mat 

@ti.kernel
def assign_partial_row() -> ti.types.matrix(3, 4, ti.i32):
    mat = ti.Matrix([[0, 0, 0, 0] for _ in range(3)])
    mat[1, 1:3] = ti.Vector([1, 2]) 
    return mat 

@ti.kernel
def augassign_rows() -> ti.types.matrix(3, 4, ti.i32):
    mat = ti.Matrix([[1, 1, 1, 1] for _ in range(3)])
    rows = ti.Matrix([[1, 2, 3, 4] for _ in range(2)])
    mat[:2, :] += rows
    return mat 

assert (assign_col() == ti.Matrix([[1, 0, 0, 0], [2, 0, 0, 0], 
                                   [3, 0, 0, 0]])).all()
assert (assign_partial_row() == ti.Matrix([[0, 0, 0, 0], [0, 1, 2, 0], 
                                           [0, 0, 0, 0]])).all()
assert (augassign_rows() == ti.Matrix([[2, 3, 4, 5], [2, 3, 4, 5], 
                                       [1, 1, 1, 1]])).all()

对于vstack(a, b)这种concatenate操作来说,你可能需要自己造一个新的Matrix然后把a, b分别拷贝到对应的位置,比如:

import taichi as ti

ti.init(arch=ti.cpu)

a = ti.Vector([1, 2, 3], ti.f32)
b = ti.Vector([3, 4, 5], ti.f32)

@ti.kernel
def concatenate(a: ti.types.vector(3, ti.f32), b: ti.types.vector(3, ti.f32)) -> ti.types.matrix(2, 3, ti.f32):
    x = ti.Matrix([[0,0,0], [0,0,0]], ti.f32)
    x[0, :] = a 
    x[1, :] = b 

    return x

x = concatenate(a, b)
print(x)