Numpy dispatch of matmul doesn't work

• Conditions OS : Linux-4.19.112±x86_64-with-Ubuntu-18.04-bionic CuPy Version : 8.3.0 NumPy Version : 1.19.5 SciPy Version : 1.4.1 Cython Build Version : 0.29.21 CUDA Root : /usr/local/cuda CUDA Build Version : 10010 CUDA Driver Version : 10010 CUDA Runtime Version : 10010 cuBLAS Version : 10201 cuFFT Version : 10101 cuRAND Version : 10101 cuSOLVER Version : (10, 2, 0) cuSPARSE Version : 10300 NVRTC Version : (10, 1) Thrust Version : 100906 CUB Build Version : 100800 cuDNN Build Version : None cuDNN Version : None NCCL Build Version : 2708 NCCL Runtime Version : 2708 cuTENSOR Version : None Device 0 Name : Tesla T4 Device 0 Compute Capability : 75
• Code to reproduce

import numpy as np
import cupy as cp

a = np.random.randn(100, 100)
a_cpu = np.asarray(a)
a_gpu = cp.asarray(a)

# numpy on CPU arrays works:
np.matmul(a_cpu, a_cpu)
# direct cupy matmul works:
cp.matmul(a_gpu, a_gpu)
# some ufunc work fine:
np.dot(a_gpu, a_gpu)

# np dispatch of matmul doesn't work:
np.matmul(a_gpu, a_gpu)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-fecacee25691> in <module>()
----> 1 np.matmul(a_gpu, a_gpu)

TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(<ufunc 'matmul'>, '__call__', array([[-1.41473897, -0.81697767,  0.93869847, ..., -0.27356864,
        -0.18850912,  1.03690492],
       [ 1.30736764, -1.17802181, -0.53378097, ..., -1.37164886,
        -1.40653998, -0.69411561],
       [-0.54080081,  0.63547369, -0.22747362, ..., -1.54336986,
        -0.2687723 ,  0.04278493],
       ...,
       [ 0.31192026,  0.82086513, -0.00734687, ...,  0.17736056,
        -1.13869101, -0.61541535],
       [-1.63717165, -2.26659611, -0.46650598, ...,  1.64617043,
        -2.28728106, -1.6082921 ],
       [ 0.4841294 , -0.1928637 , -0.66341609, ..., -1.72746173,
        -0.37958532,  0.79497399]]), array([[-1.41473897, -0.81697767,  0.93869847, ..., -0.27356864,
        -0.18850912,  1.03690492],
       [ 1.30736764, -1.17802181, -0.53378097, ..., -1.37164886,
        -1.40653998, -0.69411561],
       [-0.54080081,  0.63547369, -0.22747362, ..., -1.54336986,
        -0.2687723 ,  0.04278493],
       ...,
       [ 0.31192026,  0.82086513, -0.00734687, ...,  0.17736056,
        -1.13869101, -0.61541535],
       [-1.63717165, -2.26659611, -0.46650598, ...,  1.64617043,
        -2.28728106, -1.6082921 ],
       [ 0.4841294 , -0.1928637 , -0.66341609, ..., -1.72746173,
        -0.37958532,  0.79497399]])): 'ndarray', 'ndarray'

[emcastillo]

I think the issue here is that numpy dispatcher seems to behave differently with regular functions and ufuncs. numpy.linalg.inv is declared as a function although internally it uses an ufunc

https://github.com/numpy/numpy/blob/5c7002926ec8ad12785b36024f9da80b7017ad1c/numpy/linalg/linalg.py#L476 This uses the __array_function__ protocol.

But for matmul, it is an ufunc object, and in such objects, __array_ufunc__ seems to be dispatched a bit differently: This uses the __array_ufunc__ protocol. Here the cupy __array_ufunc__ protocol will look for the ufunc declared in CuPy but since there is no such object, it will fail, and the dispatch won’t be done.

There is a problem where cupy matmul does not support kwargs such as axes, order, etc. So mapping the function is complicated.

[emcastillo]

CuPy does not support generalized ufuncs yet. Matmul in numpy is implemented as such, so when trying to dispatch it, there is no corresponding function in CuPy. 😥