Taichiでランク多態性の関数の書き方

ランク多態性とは

さまざまな階数のテンソルを処理できる性質を私はランク多態性と呼んでいます。

なお、正式な用語ではないと思います。
こちらのスレッドの発言をそのまま使わせてもらっています。

einopsの例

たとえば以下のeinopsコードはランク多態性があります。

rearrange(x, "b ... -> b (...)")

このコードは、テンソルxの形状を以下のように変更します。
(b, d1, d2, ..., dn)　→　(b, d1 * d2 * ... * dn)
※ b：バッチサイズ、d1、d2、...、dn：他の次元
例：x:(10, 3, 4, 5)　→　(10, 60)

明示したb軸以外はどのような形状でも処理できる柔軟性があります。

Taichiとは

Taichiは、GPU並列処理をPythonで記述できるライブラリです。コンピューターグラフィックスでの利用が想定されています。

Taichiでランク多態性の関数の書き方

以下のようにすると、さまざまな形状のテンソルを処理できます。
upper_dims = ti.static(t.shape[:-2])がポイントです。
t.shapeはti.staticに記述することで、Taichiコードのコンパイル時にPythonスコープで処理しています。
なお、Taichiスコープで処理すると、upper_dimsがコンパイル時に定数ではなくなります。すると、ti.ndrangeには実行時変数を渡すことができないのでエラーになります。

参考　言語リファレンス　https://docs.taichi-lang.org/docs/language_reference

import taichi as ti
ti.init()

@ti.kernel
def tmp(t:ti.template()):
    "t : tensor with shape of [... a b]"
    print("case: rank", len(t.shape))
    upper_dims = ti.static(t.shape[:-2])
    a,b = ti.static(t.shape[-2:])
    for I in ti.static(ti.ndrange(*upper_dims)):  # unrolled serial loop
        for i,j in ti.ndrange(a,b):  # parallel loop
            t[*I,i,j] = ti.random()
            print(*I, i, j, ":", t[*I,i,j])

t = ti.field(ti.f32, shape=(2,2,2))
tmp(t) #3
t = ti.field(ti.f32, shape=(2,2,2,2))
tmp(t) #4

出力例

case: rank 3
0 0 0 : 0.443271
0 0 1 : 0.658814
0 1 0 : 0.749135
0 1 1 : 0.650376
1 0 0 : 0.948321
1 0 1 : 0.177689
1 1 0 : 0.346885
1 1 1 : 0.772146
case: rank 4
0 0 0 0 : 0.476790
0 0 0 1 : 0.989344
0 0 1 0 : 0.668613
0 0 1 1 : 0.716511
0 1 0 0 : 0.748498
0 1 0 1 : 0.891527
0 1 1 0 : 0.252390
0 1 1 1 : 0.172357
1 0 0 0 : 0.697397
1 0 0 1 : 0.212697
1 0 1 0 : 0.038212
1 0 1 1 : 0.613680
1 1 0 0 : 0.370585
1 1 0 1 : 0.226127
1 1 1 0 : 0.529289
1 1 1 1 : 0.482478