pythonnumpynumbajit

Error using jit with numpy linspace function


I am trying to use jit to compile my code. At the moment, the only library I'm using besides Numba is Numpy. When specifying nopython = True, it throws a bunch of errors with linspace arrays. I have recreated the error with a very simple function:

@jit(nopython=True)
def func(Nt):
    time = np.linspace(np.int_(0),np.int_(Nt-1),np.int_(Nt),dtype=np.int_)
    return time

Nt = 10
func(Nt)

When running, the following error message (see attached image) is shown.

Numba error

I have tried many different permutations messing with the types of the arguments with no success. All I am trying to do is make a linspace array of integers from 0 to Nt - 1. Any suggestions?


Solution

  • Remove the dtype= parameter:

    import numba
    
    
    @numba.jit(nopython=True)
    def func(Nt):
        time = np.linspace(np.int_(0), np.int_(Nt - 1), np.int_(Nt))
        return time
    
    
    Nt = 10
    print(func(Nt))
    

    Prints:

    [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]
    

    If you want integers, just recast it afterwards:

    @numba.jit(nopython=True)
    def func(Nt):
        time = np.linspace(np.int_(0), np.int_(Nt - 1), np.int_(Nt)).astype("int_")
        return time