`einsum` can't sum ellipsis axes

Description

In numpy, when optimized is set to True, einsum can sum ellipses (…). However, it seems cupy einsum doesn’t have an optimized parameter, and it can’t sum ellipses.

To Reproduce

an einsum calculation in Numpy behaves as expected:

import numpy as xp
grad=xp.random.randn(5,2,8)
x=xp.random.randn(5,2,8)
y=xp.einsum('...j,...k->jk', grad, x, optimize=True)
print(y.shape) # --> (8, 8)

However, in Cupy it produces error:

import cupy as xp
grad=xp.random.randn(5,2,8)
x=xp.random.randn(5,2,8)
y=xp.einsum('...j,...k->jk', grad, x, optimize=True)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_15756/3125293176.py in <module>
      2 grad=xp.random.randn(5,2,8)
      3 x=xp.random.randn(5,2,8)
----> 4 y=xp.einsum('...j,...k->jk', grad, x, optimize=True)

~\anaconda3\lib\site-packages\cupy\linalg\_einsum.py in einsum(*operands, **kwargs)
    527         if not options['sum_ellipsis']:
    528             if '@' not in output_subscript and -1 in dimension_dict:
--> 529                 raise ValueError(
    530                     'output has more dimensions than subscripts '
    531                     'given in einstein sum, but no \'...\' ellipsis '

ValueError: output has more dimensions than subscripts given in einstein sum, but no '...' ellipsis provided to broadcast the extra dimensions.

Installation

Wheel (pip install cupy-***)

Environment

OS                           : Windows-10-10.0.19044-SP0
Python Version               : 3.9.7
CuPy Version                 : 10.3.1
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.22.3
SciPy Version                : 1.7.1
Cython Build Version         : 0.29.28
Cython Runtime Version       : 0.29.24
CUDA Root                    : C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5
nvcc PATH                    : C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.5\bin\nvcc.EXE
CUDA Build Version           : 11050
CUDA Driver Version          : 11050
CUDA Runtime Version         : 11030
cuBLAS Version               : (available)
cuFFT Version                : 10402
cuRAND Version               : 10204
cuSOLVER Version             : (11, 1, 1)
cuSPARSE Version             : (available)
NVRTC Version                : (11, 3)
Thrust Version               : 101301
CUB Build Version            : 101301
Jitify Build Version         : 4a37de0
cuDNN Build Version          : 8302
cuDNN Version                : 8200
NCCL Build Version           : None
NCCL Runtime Version         : None
cuTENSOR Version             : None
cuSPARSELt Build Version     : None
Device 0 Name                : NVIDIA GeForce GTX 965M
Device 0 Compute Capability  : 52
Device 0 PCI Bus ID          : 0000:01:00.0

Additional Information

No response

[toslunar]

I left 'sum_ellipsis' in order to see how https://github.com/numpy/numpy/issues/10926 would be resolved.

Now it might be better allowing to sum over ellipsis axes regardless of the optimize option because

• the issue is still open and
• I’m not sure if the algorithm specified by optimize for CuPy (GPU) should be the same as NumPy.

[leofang]

Well… sure… There are many ways to do the same math problem right? 😉 If the goal here is to stay within einsum as much as possible, perhaps you can also do this:

y=xp.einsum('...j,...k->...jk', grad, x)
z=xp.sum(y, axis=(0,1))

at the cost of storing a lager intermediate array (y).