Dispatch ufunc.at methods

[dcherian]

Description

np.ufunc.at does not dispatch to cupy.ufunc.at . See below.

To Reproduce

a = cp.array([1, 2, 1, 3])
idx = np.array([0, 0, 1, 1], like=a)
out = np.full((2,), fill_value=0., like=a)
cp.add.at(out, idx, a) # works 
print(out)
np.add.at(out, idx, a)  # error
print(out)

[3. 4.]

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [59], in <cell line: 6>()
      4 cp.add.at(out, idx, a)
      5 print(out)
----> 6 np.add.at(out, idx, a)
      7 print(out)

TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(<ufunc 'add'>, 'at', array([3., 4.]), array([0, 0, 1, 1]), array([1, 2, 1, 3])): 'ndarray', 'ndarray', 'ndarray'

Installation

Conda-Forge (conda install ...)

Environment

OS                           : Linux-3.10.0-1127.18.2.el7.x86_64-x86_64-with-glibc2.10
Python Version               : 3.8.13
CuPy Version                 : 12.0.0
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.23.1
SciPy Version                : None
Cython Build Version         : 0.29.34
Cython Runtime Version       : None
CUDA Root                    : /glade/u/home/dcherian/miniconda3/envs/kvikio_env
nvcc PATH                    : None
CUDA Build Version           : 11020
CUDA Driver Version          : 11060
CUDA Runtime Version         : 11040
cuBLAS Version               : (available)
cuFFT Version                : 10502
cuRAND Version               : 10205
cuSOLVER Version             : (11, 2, 0)
cuSPARSE Version             : (available)
NVRTC Version                : (11, 4)
Thrust Version               : 101000
CUB Build Version            : 101000
Jitify Build Version         : 22da8f9
cuDNN Build Version          : None
cuDNN Version                : None
NCCL Build Version           : None
NCCL Runtime Version         : None
cuTENSOR Version             : None
cuSPARSELt Build Version     : None
Device 0 Name                : Quadro GP100
Device 0 Compute Capability  : 60
Device 0 PCI Bus ID          : 0000:18:00.0

I see similar behaviour for maximum.at and minimum.at

Additional Information

No response

[kmaehashi]

Thanks for reporting @dcherian, verified on my side.
__array_ufunc__ had to be implemented.

cupy/cupy/_core/core.pyx

Lines 1699 to 1703 in c836b93

	# Don't use for now, interface uncertain
	# elif method =='at' and name == 'add':
	# the only ufunc attribute currently
	# http://docs.cupy.dev/en/stable/reference/ufunc.html#ufunc-at
	# self.scatter_add(*inputs, **kwargs)

cc/ @asi1024

[dcherian]

Thanks for the quick fix,!

[kmaehashi]

The fix is shipped in the v12.1.0 release today. Thanks again for reporting!