如何使用 numba.njit 运行成员 function

Question

I have a class called medium and I want to increase the perfomance of a method of this class called train .

def train(self, episodeLength, T, doPlot = False, doSave = False, savePath = "./", meanTargetRewardInvAgent = None, collectiveTargetMeanReward = None):

However, this method uses the keyword self a lot, and when I call the method like this

trainCompiled = numba.njit(medium.train)
trainCompiled(episodeLength, T, doPlot = True, doSave = True, savePath = s_savePath)

it gives

Failed in nopython mode pipeline (step: analyzing bytecode)
Variable 'self' is not defined.

and when I pass medium class instance as a parameter to trainCompiled as

some keyword arguments unexpected
  File "/Volumes/Junction/Workspace/CellDev/NonEqui to ML/MLtoPhysics/AnnihilationModel/mltophysics-MxNet/main.py", line 39, in <module>
    trainCompiled(medium, episodeLength, T, doPlot = True, doSave = True, savePath = s_savePath)

So, long story short, how can I run a member function using numba.njit ?

Answer 1

You should look jit classes . "All methods of a jitclass are compiled into nopython functions". Thus, you should decorate your class medium with @jitclass .