Numba cannot resolve function (np.digitize)

Question

I get an error by numba, it is complaining it can't resolve a function.

The minimal code to reproduce my errors is

import numba
import numpy as np


@numba.jit("float64(float64)", nopython=True)
def get_cut_tight_nom(eta):
    binning_abseta = np.array([0., 10., 20.])
    
    return np.digitize(eta, binning_abseta)

I don't understand the error message

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function digitize at 0x7fdb8c11dee0>) found for signature:
 
 >>> digitize(float64, array(float64, 1d, C))
 
There are 2 candidate implementations:
      - Of which 2 did not match due to:
      Overload of function 'digitize': File: numba/np/arraymath.py: Line 3939.
        With argument(s): '(float64, array(float64, 1d, C))':
       No match.

During: resolving callee type: Function(<function digitize at 0x7fdb8c11dee0>)
During: typing of call at /tmp/ipykernel_309793/3917220133.py (8)

it seems it wants to resolve digitize(float64, array(float64, 1d, C)) and a function with the same signature is not matching?

Answer 1

It's indeed due to the signatures. The Numpy function (without Numba) of np.digitize returns an int64 (scalar or array), not a float64 as you've specified the return type of your function to be.

It seems like the Numba implementation of it requires both to always be arrays, which you'll also have to explicitly add to the signature.

So this for example works for me:

@numba.jit("int64[:](float64[:])", nopython=True)
def get_cut_tight_nom(eta):
    binning_abseta = np.array([0., 10., 20.])
    
    return np.digitize(eta, binning_abseta)

Resulting in:

But do you really need the signature in this case? Numba is able to figure it out itself as well, like:

@numba.njit
def get_cut_tight_nom(eta):
    ...

A signature can still add something if you for example want to explicitly cast float32 inputs to float64 etc.

You can also inspect what signatures Numba comes up with if you run it with some differently typed input. Running it twice with float32 & float64 as the input shows. It can help to highlight where issues like this might arise from.