numba: multiply arrays rowwise

user_11199684 0
Viewed 137

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]
    ^
numba arrays python
2 Answers

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
user_21484160

Using guvectorize is also a nice way to semi-automatically broadcast your data. A benefit is that you can compile for different targets (cpu, parallel or cuda).

For small arrays like your example, parallelization would probably only introduce overhead.

@nb.guvectorize(["void(int32[:], int32[:], int32[:])"], 
             "(n), ()->(n)", nopython=True, target="cpu")
def fct(a, b, out):
    out[:] = a * b

A = fct(a, b)

Recent Numba versions can also infer the datatypes automatically, if you're willing to provide the output array as well, and only compile for the cpu target, so:

@nb.guvectorize("(n),()->(n)", nopython=True, target="cpu")
def fct(a, b, out):
    out[:] = a * b

A = np.empty_like(a)
fct(a, b, A)
user_17554320
Related