mix data type inputs for numba njit

Question

I have a large array for operation, for example, matrix transpose. numba is much faster:

#test_transpose.py
import numpy as np
import numba as nb
import time


@nb.njit('float64[:,:](float64[:,:])', parallel=True)
def transpose(x):
    r, c = x.shape
    x2 = np.zeros((c, r))
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
   return x2


if __name__ == "__main__":
    x = np.random.randn(int(3e6), 50)
    t = time.time()
    x = x.transpose().copy()
    print(f"numpy transpose: {round(time.time() - t, 4)} secs")

    x = np.random.randn(int(3e6), 50)
    t = time.time()
    x = transpose(x)
    print(f"numba paralleled transpose: {round(time.time() - t, 4)} secs")

Run in Windows command prompt

D:\data\test>python test_transpose.py
numpy transpose: 2.0961 secs
numba paralleled transpose: 0.8584 secs

However, I want to input another large matrix, which are integers, using x as

x = np.random.randint(int(3e6), size=(int(3e6), 50), dtype=np.int64)

Exception is raised as

Traceback (most recent call last):
  File "test_transpose.py", line 39, in <module>
    x = transpose(x)
  File "C:\Program Files\Python38\lib\site-packages\numba\core\dispatcher.py", line 703, in _explain_matching_error
    raise TypeError(msg)
TypeError: No matching definition for argument type(s) array(int64, 2d, C)

It does not recognize the input data matrix as integer. If I release the data type check for the integer matrix as

@nb.njit(parallel=True) # 'float64[:,:](float64[:,:])'
def transpose(x):
    r, c = x.shape
    x2 = np.zeros((c, r))
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
    return x2

It is slower:

D:\Data\test>python test_transpose.py
numba paralleled transpose: 1.6653 secs

Using @nb.njit('int64[:,:](int64[:,:])', parallel=True) for the integer data matrix is faster, as expected.

So, how can I still allow mixed data type intputs but keep the speed, instead of creating functions each for different types?

score 1 · Accepted Answer · answered Mar 02 '22 at 21:44

So, how can I still allow mixed data type intputs but keep the speed, instead of creating functions each for different types?

The problem is that the Numba function is defined only for float64 types and not int64. The specification of the types is required because Numba compile the Python code to a native code with well-defined types. You can add multiple signatures to a Numba function:

@nb.njit(['float64[:,:](float64[:,:])', 'int64[:,:](int64[:,:])'], parallel=True)
def transpose(x):
    r, c = x.shape
    # Specifying the dtype is very important here.
    # This is a good habit to take to avoid numerical issues and 
    # slower performance in Numpy too.
    x2 = np.zeros((c, r), dtype=x.dtype)
    for i in nb.prange(c):
        for j in range(r):
            x2[i, j] = x[j][i]
   return x2

It is slower

This is because of lazy compilation. The first execution include the compilation time. THis is not the case when the signature is specified because of eager compilation is used instead.

numba is much faster

Well, not to much here considering many cores are used. In fact, the naive transposition is very inefficient on big matrices (is wast about 90% of the memory throughput in this case on large arrays). There are faster algorithms. For more information, please read this post (it only consider in-place 2D square transposition which is much simpler but the idea is the same). Also note that the wider the type, the bigger the array. The bigger the array the slower the transposition.

mix data type inputs for numba njit

1 Answers1