Лучший способ обработки нескольких констант в контейнере с numba?

Question

EDIT: Этот вопрос сильно устарел! numba теперь поддерживает Enum и namedtuple из коробки, которые обеспечивают разумное решение для группировки констант.

---

Я делаю несколько битхиптингов на питоне и хочу ускорить его с помощью numba. Для этого у меня много постоянных целочисленных значений, которые я должен обрабатывать, возможно, хорошо читаемым образом. Я хотел бы сгруппировать их вместе с объектами, подобными перечислению, имеющими все константы в одном пространстве имен, доступными с помощью оператора attribute-get. И, конечно, мне также хотелось бы, чтобы numba понимал, что там происходит, чтобы поддерживать высокие скорости с помощью jit-компиляции. Мой первый и самый наивный попробовать на что было похоже, что:

class SomeConstantsContainer: 
    SOME_NAME = 0x1 
    SOME_OTHER_CONSTANT = 0x2 
    AND_ANOTHER_CONSTANT = 0x4 

К сожалению, когда я смотрю на аннотацию, похоже, Numba не понимая, что значения являются постоянными, и он всегда возвращается к медленному доступа к объектам на объектах python. Это то, что аннотацию говорит об этом:

# $29.2 = global(SomeConstantsContainer: <class 'constants.SomeConstantContainer'>) :: pyobject 
# $29.3 = getattr(attr=SOME_VARIABLE, value=$29.2) :: pyobject 

Я знаю, что я всегда мог возвратиться к чему-то вроде этого:

from numpy import np 
SOME_STUPID_CONSTANT = np.int64(0x1) 
ANOTHER_STUPID_CONSTANT = np.int64(0x2) 

В этом случае JIT компилятор а) не нужно искать атрибут контейнера и b) точно знает, что он должен иметь дело с простым целым числом. Просто невероятно уродливо писать так. Я мог бы жить с обозначением всех констант явно целыми числами или позволить контейнеру сделать это. Тем не менее, я действительно хотел бы сгруппировать константы в контейнерах для ясности, а скомпилированную версию jit понять синтаксис и не тратить время на некоторый медленный поиск атрибутов python для каждого использования констант. Любые лучшие идеи, как сделать второй подход более похожим на первый подход, но сохраняя высокие скорости выполнения? Есть ли какой-то контейнер enum, который понимает numba, который я только что пропустил?

Edit: используя Также новый enum контейнер не помочь:

@enum.unique 
class SomeConstantsContainer(enum.IntEnum): 
    SOME_NAME = 0x1 
    SOME_OTHER_CONSTANT = 0x2 
    AND_ANOTHER_CONSTANT = 0x4 

Это дает:

# $42.3 = global(SomeConstantsContainer: <enum 'SomeConstantsContainer'>) :: pyobject 
    # $42.4 = getattr(attr=SOME_OTHER_CONSTANT, value=$42.3) :: pyobject 

Answer 1

Другой подход, но один, который все еще имеет преимущество, содержащие переменные было бы использование захваченных переменных в завернутой функции:

def make_numba_functions(): 
    SOME_NAME = 0x1 
    SOME_OTHER_CONSTANT = 0x2 
    AND_ANOTHER_CONSTANT = 0x4 

    @jit 
    def f1(x,y): 
     useful code goes here 

    @jit 
    def f2(x,y,z): 
     some more useful code goes here 

    return f1, f2 

f1,f2 = make_numba_functions() 

Вы могли бы, очевидно, объединить это с использованием класса в качестве пространства имен и распаковать содержимое внутри внешней функции.

Когда я пытаюсь это с очень простым примером и использовать inspect_types я получаю

$0.1 = freevar(A: 10.0) :: float64 

(где A просто мой постоянный). Я подозреваю, что это то, что вы хотите. Я попытался посмотреть на сгенерированный ассемблер (os.environ['NUMBA_DUMP_ASSEMBLY']='1'), но я недостаточно разбираюсь в ассемблере, чтобы выбрать соответствующую строку и подтвердить, что он делает то, что вы хотите. Однако - он компилируется с nopython=True, что, по крайней мере, предполагает, что он работает эффективно.

Немного взлома, но тот, который работает очень хорошо, надеюсь!