不同for循环下field的调用问题

在学习有限元,想要利用taichi自动并行顶层for循环的特性并行计算单元刚度矩阵。
但是发现在顶层for循环下再对数组进行完操作,只能在顶层for循环里正确调用数组。
比如在如下代码中,在循环 for element in ti.ndrange(1):中似乎无法正确使用数组J,这是为什么?如何解决?

import taichi as ti


@ti.kernel
def test():
    for thread_i in ti.ndrange(1):
        for element in ti.ndrange(1):
            for i, j in ti.ndrange((0, 3), (0, 3)):
                J[thread_i][i, j] *= 0.
                for k in ti.ndrange(3):
                    J[thread_i][i, j] += 1.
            print("second for:", J[thread_i])    # 数组没有更新
        print("top for:", J[thread_i])    # 正确调用


J = ti.Matrix.field(3, 3, dtype=ti.f32, shape=3)
j_ = np.zeros((3, 3, 3), dtype=np.float32)
J.from_numpy(j_)
test()

输出如下:

second for: [[0.000000, 0.000000, 0.000000], [0.000000, 0.000000, 0.000000], [0.000000, 0.000000, 0.000000]]
top for: [[3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000]]

求助,另外,一般在节省内存的情况下如何使用taichi加速有限元。

似乎我的程序能正常的运行,两次print都是对的,或许是版本问题?

[Taichi] version 1.7.2, llvm 15.0.7, commit 0131dce9, osx, python 3.11.0
[Taichi] Starting on arch=arm64
second for: [[3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000]]
top for: [[3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000], [3.000000, 3.000000, 3.000000]]

好像以前用taichi的时候,在ti.init()里面加入了debug=True,反而使我的print有些问题,会不会和这个有关?