numbaに複雑な型を渡す
素直にnp.arrayにできるならそれが楽でいいのだけど、実用上は「キーが整数である辞書で、値は可変長のリスト。リストの要素は2整数のタプル」みたいな複雑な型のことがあり(実話)これをnp.arrayに詰め込もうとしてエンバグするのは不毛。
この型は、生のPythonの辞書やリストのままではnumbaでうまく扱えないのだが、typed.Dictとtyped.Listを使えば実現できる。問題はこれを関数に渡す時に関数の型シグニチャがどうなるか。
numba.typeofで調べると
DictType[int64,ListType[UniTuple(int64 x 2)]]
と表示されるが、これをそのまま型シグニチャに書くと動かない。修正点は
UniTuple(int64 x 2) → UniTuple(int64, 2)
XXXType[xxx] → XXXType(xxx)
(たぶんこれ、パーサを作らないでevalで解決してる。だからPython的文法にする必要がある)
文字列にせず、typeofで推定させる手もある
code:python
@nb.jit(nb.typeof((1.0,1.0))(nb.double),nopython=True)
def f(a):
return a,a
これで当初の目的である複雑な型の値をnumbaコンパイル済みの関数に渡すことができた。
追加: atcoderのnumbaは0.48と古いのでtyped.Listが引数を取れないようだ
code:python
import sys
import numba
def foo(x):
for k in x:
print("k:", k)
for a, b in xk:
print(a, b)
if sys.argv-1 == "-c":
# numba compile
print("compiling")
from numba.pycc import CC
cc = CC('my_module')
cc.export(
'foo', 'void(DictType(int64,ListType(UniTuple(int64,2))))')(foo)
cc.compile()
exit()
else:
x = numba.typed.Dict()
x1 = numba.typed.List((1, 2), (3, 4))
x100 = numba.typed.List((100, 200))
if sys.argv-1 == "-p":
print(numba.typeof(x))
# => DictType[int64,ListTypeUniTuple(int64 x 2)]
else:
from my_module import foo
foo(x)
numba