Add support for `np.matmul`

[stuartarchibald]

Feature request

As title. Noted missing functionality from #3803.

[mohdkashif93]

Hey, is it okay if I take this issue up ?

[stuartarchibald]

@mohdkashif93 Certainly! Thanks!

As a new contributor, some instructions for getting set up for development are here and are a great place to start.

The use of Numba's extension API @overload decorator is strongly recommended for this task, documentation is here.

It may be useful to take a look at examples of pull requests adding support for a NumPy function that have already been done if guidance is needed... here's a couple:

• np.roll: https://github.com/numba/numba/pull/3563/files
• np.roots: https://github.com/numba/numba/pull/2135/files

Thanks again!

[Quetzalcohuatl]

Note that np.dot(a, b) where a and b are both 2D arrays, it is the same as np.matmul. And since np.dot is supported, then we win. https://docs.scipy.org/doc/numpy/reference/generated/numpy.dot.html

[Quetzalcohuatl]

However, more importantly the operator @ which is matrix multiplication between numpy arrays is also supported. so just use that..

[mohdkashif93]

@Quetzalcohuatl Thanks for the heads up 👍 I will get this done quickly.

[stavoltafunzia]

Feature request

As title. Noted missing functionality from #3803.

I also vote for this feature request. Unfortunately with the operator @ it's not possible to specify a pre-allocated output array.

Edit: not needed as numpy.dot supports the output variable as argument.

[samaocarpenter]

Giving a +1 to this feature request. Having np.matmul supported in numba would be pretty great.

[njriasan]

@mohdkashif93 are you still working on this. If not I'm happy to work on this feature.

[externalsupplierstaff]

No implementation of function Function(<built-in function matmul>) found for signature:
 
 >>> matmul(array(float64, 3d, C), readonly array(float64, 2d, F))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'MatMul.generic': File: numba/core/typing/npydecl.py: Line 994.
    With argument(s): '(array(float64, 3d, C), readonly array(float64, 2d, F))':
   Rejected as the implementation raised a specific error:
     TypingError: '@' only supported on 1-D and 2-D arrays
  raised from /home/francomarchesoni/miniconda3/envs/sandbox/lib/python3.9/site-packages/numba/core/typing/npydecl.py:908

During: typing of intrinsic-call at /home/francomarchesoni/huawei/code/exploration/distort_and_recover/rgb2lab.py (26)

File "exploration/distort_and_recover/rgb2lab.py", line 26:
def jitrgb2lab(arr):
    <source elided>
    arr = flattened.reshape(arr.shape)
    arr = arr @ xyz_from_rgb_T
    ^

However, more importantly the operator @ which is matrix multiplication between numpy arrays is also supported. so just use that..

it seems that it doesn't work for 3d tensors 😢

[timhdesilva]

Any update on this? I have the following example that does not work when I use np.matmul or the @ decorator for matrix multiplication. Any guidance on how to workaround this would be great!

from numba import jit
import numpy as np

@jit(nopython = True)
def smm_loop(simdata, m_data, W, perc):
    N = np.shape(simdata)[0]
    smm_objs = np.zeros((N, 1))
    for i in range(N):
        m_model = simdata[i, :].reshape((-1,1))
        if perc:
            m_diff = np.divide(m_model - m_data, m_data)
        else:
            m_diff = m_model - m_data
        val = np.matmul(np.matmul(np.transpose(m_diff), W), m_diff)
        smm_objs[i, 0] = val
    return smm_objs


dt = np.random.rand(100,10)
m = np.random.rand(10,1)
W = np.identity(10)
perc = False

out = smm_loop(dt, m, W, perc)

[stavoltafunzia]

Why not simply using numpy.dot as already pointed out? For 2d arrays it performs matrix multiplication.

[timhdesilva]

@stavoltafunzia - Thanks for the response! When attempting to do so with the following code:

from numba import jit
import numpy as np

@jit(nopython = True)
def smm_loop(simdata, m_data, W, perc):
    N = np.shape(simdata)[0]
    smm_objs = np.zeros((N, 1))
    for i in range(N):
        m_model = simdata[i, :].reshape((-1,1))
        if perc:
            m_diff = np.divide(m_model - m_data, m_data)
        else:
            m_diff = m_model - m_data
        # val = np.matmul(np.matmul(np.transpose(m_diff), W), m_diff)
        val = np.dot(np.dot(np.transpose(m_diff), W), m_diff)
        smm_objs[i, 0] = val
    return smm_objs


dt = np.random.rand(100,10)
m = np.random.rand(10,1)
W = np.identity(10)
perc = False

out = smm_loop(dt, m, W, perc)

I get the following error:

Traceback (most recent call last):

  File "<ipython-input-1-c37f901b015f>", line 25, in <module>
    out = smm_loop(dt, m, W, perc)

  File "/Users/timdesilva/opt/anaconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 420, in _compile_for_args
    error_rewrite(e, 'typing')

  File "/Users/timdesilva/opt/anaconda3/lib/python3.8/site-packages/numba/core/dispatcher.py", line 361, in error_rewrite
    raise e.with_traceback(None)

TypingError: No implementation of function Function(<built-in function setitem>) found for signature:
 
setitem(array(float64, 2d, C), Tuple(int64, Literal[int](0)), array(float64, 2d, C))
 
There are 16 candidate implementations:
  - Of which 16 did not match due to:
  Overload of function 'setitem': File: <numerous>: Line N/A.
    With argument(s): '(array(float64, 2d, C), UniTuple(int64 x 2), array(float64, 2d, C))':
   No match.

During: typing of setitem at <ipython-input-1-c37f901b015f> (16)