move out.ndim checks to be prior to the CUB calls
@@ -46,11 +46,16 @@ cdef extern from 'cupy_cub.h':
def reduce_sum(core.ndarray x, out=None, bint keepdims=False):
cdef core.ndarray y
cdef core.ndarray ws
- cdef int dtype_id
+ cdef int dtype_id, ndim_out
cdef size_t ws_size
cdef void *x_ptr
cdef void *y_ptr
cdef void *ws_ptr
+ ndim_out = keepdims
+ if out is not None and out.ndim != ndim_out:
+ raise ValueError(
+ "output parameter for reduction operation sum has the wrong "
+ "number of dimensions")
x = core.ascontiguousarray(x)
y = core.ndarray((), x.dtype)
x_ptr = <void *>x.data.ptr
@@ -63,10 +68,6 @@ def reduce_sum(core.ndarray x, out=None, bint keepdims=False):
if keepdims:
y = y.reshape((1,))
if out is not None:
- if y.ndim != out.ndim:
- raise ValueError(
- "output parameter for reduction operation sum has the wrong "
- "number of dimensions")
out[...] = y
y = out
return y
@@ -100,11 +101,16 @@ def can_use_reduce_sum(x_dtype, Py_ssize_t ndim, dtype=None, axis=None):
def reduce_min(core.ndarray x, out=None, bint keepdims=False):
cdef core.ndarray y
cdef core.ndarray ws
- cdef int dtype_id
+ cdef int dtype_id, ndim_out
cdef size_t ws_size
cdef void *x_ptr
cdef void *y_ptr
cdef void *ws_ptr
+ ndim_out = keepdims
+ if out is not None and out.ndim != ndim_out:
+ raise ValueError(
+ "output parameter for reduction operation sum has the wrong "
+ "number of dimensions")
x = core.ascontiguousarray(x)
y = core.ndarray((), x.dtype)
x_ptr = <void *>x.data.ptr
@@ -117,10 +123,6 @@ def reduce_min(core.ndarray x, out=None, bint keepdims=False):
if keepdims:
y = y.reshape((1,))
if out is not None:
- if y.ndim != out.ndim:
- raise ValueError(
- "output parameter for reduction operation min has the wrong "
- "number of dimensions")
out[...] = y
y = out
return y
@@ -141,11 +143,16 @@ def can_use_reduce_min(x_dtype, Py_ssize_t ndim, dtype=None, axis=None):
def reduce_max(core.ndarray x, out=None, bint keepdims=False):
cdef core.ndarray y
cdef core.ndarray ws
- cdef int dtype_id
+ cdef int dtype_id, ndim_out
cdef size_t ws_size
cdef void *x_ptr
cdef void *y_ptr
cdef void *ws_ptr
+ ndim_out = keepdims
+ if out is not None and out.ndim != ndim_out:
+ raise ValueError(
+ "output parameter for reduction operation sum has the wrong "
+ "number of dimensions")
x = core.ascontiguousarray(x)
y = core.ndarray((), x.dtype)
x_ptr = <void *>x.data.ptr
@@ -158,10 +165,6 @@ def reduce_max(core.ndarray x, out=None, bint keepdims=False):
if keepdims:
y = y.reshape((1,))
if out is not None:
- if y.ndim != out.ndim:
- raise ValueError(
- "output parameter for reduction operation max has the wrong "
- "number of dimensions")
out[...] = y
y = out
return y
