はじめに

NumbaではNumba化関数をある程度まで関数オブジェクトとして扱うことができる。
一方でその機能は完全でないとも書かれている [公式]
どこまでできるか検証した。

検証環境:

Python 3.11.4 + Numba 0.58.1
Python 3.11.4 + Numba 0.57.0

目次

Numba化関数は別のNumba化関数オブジェクトを引数として受け取り、動作を切り替えることができる
Numba化関数オブジェクトを返り値で返す
関数オブジェクトの型指定方法
おわり

Numba化関数は別のNumba化関数オブジェクトを引数として受け取り、動作を切り替えることができる

下記が動作する。

import numba
from numba import jit
print("Numba:", numba.__version__)

@jit(nopython=True, cache=True)
def add(a, b):
    return a + b

@jit(nopython=True, cache=True)
def sub(a, b):
    return a - b


@jit(nopython=True, cache=True)
def receive_f(func, n):
    return func(10, n)

print(receive_f(add, 2))
print(receive_f(sub, 2))

Numba化関数オブジェクトを返り値で返す

Numba化関数を修飾して別のNumba化関数を返すデコレータを作りたい

加工せずに返すことはできる

下記は動作した。(使い道は思いつかない)

@jit(nopython=True, cache=True)
def f_passer(func):
    return func

passed = f_passer(add)
print(passed(10, 3))

関数を加工して関数オブジェクトを返す関数 (デコレータ) - 失敗例

Pythonでのデコレータのコードをそのまま @jit してもエラーになる。
公式にNumba化関数内の関数は引数として渡したり返り値として返したりできないと書かれている。

# これはエラー
@jit(nopython=True, cache=True)
def mydecorator1(func):
    def wrapper(arg1, arg2):
        print("begin")
        ret = func(arg1, arg2)
        print("finished")
        return ret
    return wrapper

decorated_add1 = mydecorator1(add)
# Error

関数を加工して関数オブジェクトを返す関数 (デコレータ) - 成功例

外側関数をNumba化せず、返り値にする内側関数に対して @jit すればデコレータとして動作した。

# ここには@jitしない"
def mydecorator2(func):
    # wrapperを@jitする
    @jit(nopython=True, cache=True)
    def wrapper(*args):
        print("begin")
        ret = func(*args)
        print("finish")
        return ret
    return wrapper

# 関数呼び出し記法でNumba化関数addを修飾する
decorated_add2 = mydecorator2(add)

# ＠デコレータ記法で関数定義時に修飾する
@mydecorator2
@jit(nopython=True, cache=True)
def decorated_mul(a, b):
    return a * b

print("修飾された関数をPython領域から呼び出す")
print(decorated_add2(100, 2))
print(decorated_mul(100, 2))

print("修飾された関数はNumba化関数なのでNumba化関数内から呼び出せる")
@jit(nopython=True, cache=True)
def call_decorated_f():
    res1 = decorated_add2(100, 3)
    print(res1)
    res2 = decorated_mul(100, 3)
    print(res2)
    return None
call_decorated_f()

print("修飾された関数はNumba化関数なのでNumba化関数へ引数で渡せる")
@jit(nopython=True, cache=True)
def receive_decorated_f(func):
    res = func(100, 4)
    return res
print(receive_decorated_f(decorated_add2))
print(receive_decorated_f(decorated_mul))

デコレータ自体がNumba化されないのが少し気になったが、修飾処理そのものを高速化する必要がある状況はそうそうないはず。
そもそも @jit デコレータ自体がPython領域で動くものなのでこれ以上の方法はなさそう。

クロージャ定数のキャプチャは型次第だができた

クロージャとしてキャプチャした値を定数として利用することはできた。
「Numba化関数内から外側のPython領域へのアクセス」となるのでやや不安。
なお、文字列や数値やNumpy配列はキャプチャできたが listやdictはできなかった。可能な型と不可能な型があると思われる。

def mydecorator3(func):
    NAME = func.__name__
    @jit(nopython=True, cache=True)
    def wrapper(*args):
        # クロージャとして定数NAMEをキャプチャさせる
        print("begin", NAME)
        ret = func(*args)
        print("finish", NAME)
        return ret
    return wrapper

decorated_add3 = mydecorator3(add)
print(decorated_add3(100, 2))
# "begin add", "finish add"

クロージャでキャプチャした変数の変更は失敗した

同様にキャプチャした変数の更新を試みたが、プログラムは動いたものの値は更新されなかった。
「Numba化関数内から外側のPython領域へのアクセス」となるのでコンパイル時に凍結されたと思われる。
参考: Global and closure variables

def mydecorator4(func):
    count = 0
    @jit(nopython=True, cache=True)
    def wrapper(*args):
        # クロージャに変数countをキャプチャさせ、呼び出し回数を記録させたい
        nonlocal count
        count += 1
        print("begin,", "count:", count)
        ret = func(*args)
        print("finish")
        return ret
    return wrapper

decorated_add4 = mydecorator4(add)
print(decorated_add4(1000, 1))
print(decorated_add4(1000, 1))
print(decorated_add4(1000, 1))
# ! 動作するがcountの変更が反映されない!

上記はint型での結果だが、キャプチャ対象の型によってはコンパイル時エラーとなった。

キャプチャした変数の変更を保持できないのでデコレータとしての応用範囲は限られそう。

関数オブジェクトの型指定方法

型指定しない

これまでいろいろな条件でNumbaの速度測定をしたが、 @jit に型指定してもしない場合と処理速度は変わらなかったので、明確な必要性がない限り@jitへの型指定はしない方がよいと考えている。
jitclassやAOTでは型指定が必須なので後述の方法をとる。

また、型指定しない場合、処理内容の切り替えをしても動作可能な組み合わせなら動いてくれる。(最初のコード参照)

定義済み関数から型を得られるがその場合は関数オブジェクトの入れ替えに非対応になる

定義済みのNumba化関数の名前を nmbfunc としたとき、次のいずれかでその型指定オブジェクトが得られた。

numba.typeof(nmbfunc)
nmbfunc._type
nmbfunc._numba_type_
これを使って次のjitclassは動いた。
ただし、この型指定オブジェクトは各関数毎に異なる型として扱われる。つまり、変数型を型指定する際にこれらを与えると、処理内容の切り替え (本記事の add と sub を入れ替えて使うようなこと) ができない。関数の入れ替えに非対応なので使いどころが限られる。

import numba
from numba import jit, int64
from numba.experimental import jitclass

@jit(nopython=True, cache=True)
def add(a, b):
    return a + b

@jit(nopython=True, cache=True)
def sub(a, b):
    return a - b


@jitclass([
    ("arg1", int64),
    ("func", numba.typeof(add)),  # add関数の型
    ])
class MyClass():
    def __init__(self, arg1, f):
        self.arg1 = arg1
        self.func = f
    def set_func(self, g):
        self.func = g
    def calc(self, arg2):
        return self.func(self.arg1, arg2)

@jit(nopython=True, cache=True)
def use_myclass(arg1, arg2):
    ins = MyClass(arg1, add)
    print(ins.calc(arg2))
    # ↓ このコメントアウトを解除するとエラーになる (self.funcをaddからsubに入れ替えできない)
    # ins.set_func(sub)
    print(ins.calc(arg2))

use_myclass(10, 3)

後述の入れ替え対応の型指定に対するこちらの優位点は、型指定が楽なこと。
また、こちらは同じ関数オブジェクトに異なる引数型を変えても追加コンパイルして動くという性質がある。

関数オブジェクトの入れ替えに対応する型指定 : `numba.types.FunctionType(返り値の型(引数の型))`

関数オブジェクトを入出力型で抽象化した場合の型指定は numba.types.FunctionType(返り値の型(引数の型)) で得られる。
こちらは関数オブジェクトの入れ替えに対応する。
なお、Numbaのソースコードを見るとこのオブジェクトが登場したのはNumba 0.49.0 以降だったため、古いNumbaでは使えない模様。

from numba import jit, int64, types
from numba.experimental import jitclass

@jit(nopython=True, cache=True)
def add(a, b):
    return a + b

@jit(nopython=True, cache=True)
def sub(a, b):
    return a - b


@jitclass([
    ("arg1", int64),
    ("func", types.FunctionType(int64(int64, int64))),  # 型指定
    ])
class MyClass():
    def __init__(self, arg1, f):
        self.arg1 = arg1
        self.func = f
    def set_func(self, g):
        self.func = g
    def calc(self, arg2):
        return self.func(self.arg1, arg2)

@jit(nopython=True, cache=True)
def use_myclass(arg1, arg2):
    ins = MyClass(arg1, add) # self.func = add として初期化
    print(ins.calc(arg2))    # 保持させた add を使用
    ins.set_func(sub)        # self.func を sub に入れ替え
    print(ins.calc(arg2))    # 保持させた sub を使用

use_myclass(10, 3)

おわり

Numba化関数はオブジェクトとして引数に渡すことができ、処理内容の切り換えに利用できる。
一応デコレータも作れたが、クロージャ変数のキャプチャが限られた型のみかつ変更不可なので応用は限られそう。