How to declare an array of enum values that is compatible with numba?

Numba Support: How do I do ...?
1

I’m new to numba and trying to numbafy an existing python script to hopefully make it run faster.

I’m declaring an array that I want to pass to the constructor of my CheckerBoardState class, but it’s being rejected because it is of type list(Enum<int64>(Cell)) but it need to be type ListType[<enum 'Cell'>]. How do I specify the correct type?

Cannot cast reflected list(Enum<int64>(Cell))<iv=None> to ListType[<enum 'Cell'>]: %"inserted.parent" = insertvalue {i8*, i8*} %"inserted.meminfo.1", i8* %"arg.pieces.1", 1

from numba.experimental import jitclass
from numba import njit, types
from typing import List
from enum import Enum


class Cell(Enum):
    EMPTY = 0
    BLACK = 1
    BLACK_KING = 2
    WHITE = 3
    WHITE_KING = 4

@jitclass([('black_to_move', types.bool), ('pieces', types.ListType(Cell))])
class CheckerBoardState:
    def __init__(self, black_to_move:bool, pieces:List[Cell]):
        self.black_to_move = black_to_move
        self.pieces = pieces

board_start:CheckerBoardState = CheckerBoardState(True, [
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 

    Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, 
    Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, 

    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
])
2

You can try

from numba import typed

@jitclass([('black_to_move', types.bool), ('pieces', types.ListType(types.EnumMember(Cell, types.int64)))])
class CheckerBoardState:
    def __init__(self, black_to_move:bool, pieces:List[Cell]):
        self.black_to_move = black_to_move
        self.pieces = pieces

board_start:CheckerBoardState = CheckerBoardState(True, typed.List([
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 
    Cell.BLACK, Cell.BLACK, Cell.BLACK, Cell.BLACK, 

    Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, 
    Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, Cell.EMPTY, 

    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
    Cell.WHITE, Cell.WHITE, Cell.WHITE, Cell.WHITE, 
]))
2 Likes
3

That being said, List is relatively slow. Numpy array will likely be faster.