Modify in place a slice of numpy array

[FrankFrank9]

Hello,

I'm facing an issue where I can't modify an underlying numpy array from taichi. I would like to modify upts directly but in my original code I need to access a slice of it namely the values allocated in upts_. Here is a MWE of what I'm trying to achieve

import numpy as np
import taichi as ti

ti.init(arch=ti.cpu)

upts = np.zeros((3, 10), dtype=np.float64)


@ti.func
def taich_fun(u: ti.template()):
    u[0] = 1.0
    u[2] = 2.0


@ti.kernel
def foo(upts: ti.types.ndarray()):
    for i in range(10):
        upts_ = ti.Vector([0.0 for _ in range(3)], dt=ti.f64)
        for j in range(3):
            upts_[j] = upts[j, i]

        taich_fun(upts_)


foo(upts)
print(upts)

Do you have any idea on how to work around this?

[Ives0721]

Sorry I can't fully understand your meaning yet, so are those my following comment match what you means?

Variable upts_ in every range-for-loop of kernel foo is a newly created object, not a memoryview of upts[j, :]. If you want to modify the value of upts[j, :] with the values created by taich_fun(upts_), the code of your kernel may look like :

@ti.kernel
def foo(upts: ti.types.ndarray()):
    for i in range(10):
        upts_ = ti.Vector([0.0 for _ in range(3)], dt=ti.f64)

        taich_fun(upts_)

        for j in range(3):
            upts[j, i] = upts_[j]

And the result I run in Taichi v1.6.0 is:

[Taichi] version 1.6.0, llvm 15.0.1, commit f1c6fbbd, win, python 3.9.16
[Taichi] Starting on arch=x64
[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]]