Search code examples
pythonarraysnumba

numba: multiply arrays rowwise


I have numpy arrays shapes (2,5) and (2,) and I want to multiply them rowvise

a = np.array([[3,5,6,9,10],[4,7,8,11,12]])
b = np.array([-1,2])

From numpy: multiply arrays rowwise I know this works with numpy:

a * b[:,None] 

which gives correct output

array([[ -3,  -5,  -6,  -9, -10],
       [  8,  14,  16,  22,  24]])

But with numba it does not work any more, I get a bunch of error messages beyond me.

Code:

import numpy as np
from numba import njit

@njit()
def fct(a,b):
    c = a * b[:,None]
    return c

a = np.array([[3,5,6,9,10],[4,7,8,11,12]])
b = np.array([-1,2])
A = fct(a, b)
print(A)

I put this code in a file named numba_questionA.py. Running it gives error messages:

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function getitem>) found for signature:
 
 >>> getitem(array(int32, 1d, C), Tuple(slice<a:b>, none))
 
There are 22 candidate implementations:
  - Of which 20 did not match due to:
  Overload of function 'getitem': File: <numerous>: Line N/A.
    With argument(s): '(array(int32, 1d, C), Tuple(slice<a:b>, none))':
   No match.
  - Of which 2 did not match due to:
  Overload in function 'GetItemBuffer.generic': File: numba\core\typing\arraydecl.py: Line 162.
    With argument(s): '(array(int32, 1d, C), Tuple(slice<a:b>, none))':
   Rejected as the implementation raised a specific error:
     TypeError: unsupported array index type none in Tuple(slice<a:b>, none)
  raised from numba_questionA.py

During: typing of intrinsic-call at numba_questionA.py
During: typing of static-get-item at numba_questionA.py

File "numba_questionA.py", line 6:
def fct(a,b):
    c = a * b[:,None]
    ^

Solution

  • Numba says it can't use None as an array index, so you can replace

    b[:, None]
    

    by

    b.reshape(-1, 1)
    

    However, Numba won't likely be faster than Numpy for an expression like a * b[:,None].

    But, if your arrays are really large, you can take advantage of Numba's parallelization:

    @nb.njit(parallel=True)
    def fct(a, b):
        c = np.empty_like(a)
        for i in nb.prange(a.shape[1]):
            c[:, i] = a[:, i] * b
        return c