It appears that keepdims does not keep all of the reduction dimensions when CUB is enabled and axis is a tuple containing more than one axis.
- Conditions (you can just paste the output of
python -c 'import cupy; cupy.show_config()')
CuPy Version : 7.0.0rc1 (build from master branch)
CUDA Root : /usr/local/cuda
CUDA Build Version : 10010
CUDA Driver Version : 10010
CUDA Runtime Version : 10010
import cupy
x = cupy.ones((64, 64))
cupy.cuda.cub_enabled=False
out = x.sum(axis=(0, 1), keepdims=True)
print("CUB DISABLED shape={}".format(out.shape))
cupy.cuda.cub_enabled=True
out = x.sum(axis=(0, 1), keepdims=True)
print("CUB ENABLED shape={}".format(out.shape))gives
CUB DISABLED shape=(1, 1)
CUB ENABLED shape=(1,)
cc @leofang
It appears that keepdims does not keep all of the reduction dimensions when CUB is enabled and axis is a tuple containing more than one axis.
python -c 'import cupy; cupy.show_config()')gives
cc @leofang