Question about Numba (nopython) implementation of np.matmul

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.

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

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

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

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