How can I make this function jit-compatible?

Numba Community Support
1

Hello,

I am trying to get my head around numba. For this, I tried to implement it in one of my smaller Jupyter scripts, which has two similar functions. I would like to link one of them here, and if you can help out, I’ll try to take what I learnt and work on the other one.

This is my function:

def standard_combos(elements: Tuple, depth: int = 4, needed = Literal["elements", "points"]):

  combinations = ()

  for i in range(1,depth):
    combinations = combinations + tuple(itertools.combinations(elements,i))

  match needed:
    case "elements": return len(set(combinations))
    case "points": 
      elementcounts = itertools.chain.from_iterable(map(list,combinations))
      count = collections.Counter(elementcounts)
      return count

I keep getting this error:

2

Hey @Rieleniel ,

Numba has its own implementation of many NumPy functions, that’s why you can use them.
Unfortunately, the itertools module is not implemented in Numba, that’s why you cannot accelerate itertools functions using Numba.
You can use these functions in object mode, but this mode probably does not provide performance improvement.

import itertools
import collections
import numba

@numba.jit(nopython=False, forceobj=True)
def standard_combos(elements, depth=4, needed="elements"):
    combos = ()
    for i in range(1, depth):  # <= is this correct or should it be range(1, depth+1)
        combos = combos + tuple(itertools.combinations(elements, i))
    match needed:
        case "elements":   
            return len(set(combos))
        case "points":
            elementcounts = itertools.chain.from_iterable(map(list, combos))
            count = collections.Counter(elementcounts)
            return count