Numba: How to make a compiled function return the same dtype as the non-compiled version?

Viewed 17

I'm trying to get a very simple function to output the same dtype when compiled vs non-compiled. Currently the non-compiled version returns a np.float32, as expected, whereas the compiled version returns float.

Example:

a = np.array([1.,2.,3.], dtype=np.float32)

def np_sum(a):
    return np.sum(a)

@numba.njit()
def np_sum_numba(a):
    return np.sum(a)

type(np_sum(a)) # numpy.float32
type(np_sum_numba(a)) # float

How can I make the compiled function also return a numpy.float32 ?

0 Answers
Related