Question about Numba (nopython) implementation of np.matmul

Numba Support: How do I do ...?
1

Greetings. I have a problem involving the use of np.matmul inside a function which I am trying to accelerate with the decorator @njit. Since np.matmul causes an error when used in combination with the decorator, I have been trying to use @overload to implement it, but have run into a problem. A small example showcasing the problem is given below:

#importing the requisite libraries
import numpy as np
from numba import njit, types
from numba.extending import overload

#trying to implement a @njit version of np.matmul
@overload(np.matmul)
def jit_matmul(x,y):
    if isinstance(x,types.Array) and isinstance(y,types.Array):
        def impl(x,y):
            return x+y
    return impl

#just defining a function f(x) for convenience and future use
@njit
def f(x,y):
    return np.matmul(x,y)

#These matrices are to check whether the code is working
a = np.matrix([[2.0, 4.0], [4.0, 5.6]])
b = np.matrix([[2.1, 4.1], [4.1, 5.6]])

However, I get an error to the effect of:
`TypingError: Failed in nopython mode pipeline (step: native lowering)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<ufunc ā€˜matmulā€™>) found for signature:

matmul(float64, float64)`

There is some other stuff in the error, but I believe this is the important part. Iā€™m quite a beginner in Numba and would appreciate any assistance. Thank you.

2

Do you really need np.matmul? If so, itā€™s probably best to give an example that emphasizes this. Your current example would also work fine with the normal dot-product operator. The differences are explained in the ā€œnotesā€ at:
https://numpy.org/doc/stable/reference/generated/numpy.matmul.html

import numpy as np
from numba import njit

@njit
def f(x,y): return x @ y

a = np.array([[2.0, 4.0], [4.0, 5.6]])
b = np.array([[2.1, 4.1], [4.1, 5.6]])

r_nb = f(a, b)
r_py = a @ b
r_np = np.matmul(a, b)

np.testing.assert_array_equal(r_nb, r_py)
np.testing.assert_array_equal(r_nb, r_np)
1 Like
3

Thank you for the reply. This solves the problem rather neatly. Just one last question, if you donā€™t mind. In my original code snippet, am I doing something wrong? This is just for my curiosity, your previous reply has solved my immediate problem.

1 Like
4

No not all, what you tried is perfectly reasonable as far as Iā€™m concerned. But even with plain Numpy (no Numba) the use both np.matrix (over np.array) and using np.matmul (over @ or np.dot) are somewhat of a niche, and therefore less well supported in general, and thatā€™s also true for Numba as youā€™ve now found out.

As said, there are some subtle differences between @/np.dot and np.matmul, so there is a distinct use case for the latter, hence my question. But for most people a ā€œplainā€ dot product on two arrays is often fine, so perhaps a better default until you really need the specifics of matrix/matmul.

1 Like
5

I take your point. Thank you again for your assistance and the information.