Numba, Jitclass w/ nopython mode and dictionaries

2.8k Views Asked by At

trying to figure out how to skip a class method while using jitclass.

Have a pretty big recursive model (pretty much, a massive for-loop), which - given path-dependent calculations, cannot be vectorized with straight Numpy.

The class runs through a series of numpy arrays, with generally numba-friendly syntax, however I have one section which calls a few of the methods in an ordered fashion:

def operations(self, i, ops_order_config):

    ops_dict = self.ops_dict

    for index in range(len(waterfall_config)):
        try:
            if isinstance(ops_config[index], tuple):
                ops_dict[ops_config[index][0]](i, ops_config[index][1])
            else:
                ops_dict[ops_config[index]](i)
        except KeyError:
            pass

This part of the model is pretty crucial for flexibility - the "config" is an ordered list of tuples which contain the appropriate method to call, and the respective parameters. The ops_dict holds the actual self. that is called from the config, with proper parameters.

If I'm making a jitclass, is there any way to just skip over jitting this dictionary aspect?

2

There are 2 best solutions below

3
On BEST ANSWER

No, if you make a jitclass every attribute has to be typed and dictionaries or lists/tuples containing functions (even if jitted) aren't supported as of numba 0.34. For example trying to use dict or object as type:

import numpy as np
from numba import jitclass

spec = [('dct', dict)]

@jitclass(spec)
class ClsWithObject(object):
    def __init__(self, value):
        self.dct = {}

Throws a TypeError:

TypeError: spec values should be Numba type instances, got <class 'dict'>

Besides, using isinstance as well as try and except don't work in nopython-mode neither.

Your best option would be to use a jitted function that is called from within a pure Python class.

1
On

With regards to the use of Dictionaries in Numba compiled functions, as MSeifert said, Numba does not support this. Within my own work I've run into this problem, and found an implementation (not created by me) of dictionaries in Numba that worked great, it's GitHub repository can be found here.