Why won't Numba let me define a variable?

[transversalist]

I have this piece of code I want to run faster, but when It defines “num”, it gives an odd error message

[sklam]

Can you show more context? There are usually more details in the traceback.
My guess is that the problem is somewhere in the expression on the RHS of =.

[transversalist]

I’m fairly confident the issue is the prime function from sympy, but I don’t know how to fix it.

[milton]

Welcome to numba!

Objects used in jit’ed context should have well-defined numba types.

Python function is not one of them. In particular, sympy.prime is simply an instance of Python function, so it cannot be typed within a numba jit-compiled function.

If you must use sympy.prime within a jit-compiled function, you can overload it:

from numba import njit
from numba.core.types import Integer
from numba.extending import overload


from sympy import prime


def get_nth_prime(n):
    """
    Your favorite algorithm goes here, but make sure it is numba-friendly.
    """


@overload(prime, strict=False)
def ol_prime(n_ty):
    if not isinstance(n_ty, Integer):
        raise TypeError(f"Expected integer, got {n_ty}")
    return get_nth_prime


@njit
def your_function(n):
    p = prime(n) # now you can use it, thanks to overload
    ...