Fix reduction contiguous size
@@ -183,7 +183,7 @@ cdef shape_t _set_permuted_args(
cdef Py_ssize_t _get_contiguous_size(
- list args, tuple params, Py_ssize_t ndim,
+ list args, tuple params, list out_shape, Py_ssize_t ndim,
Py_ssize_t out_ndim) except -1:
'''
get contiguous size in the *output* axis (not *reduce* axis!)
@@ -203,7 +203,7 @@ cdef Py_ssize_t _get_contiguous_size(
for j in range(out_ndim):
if a._strides[ndim-j-1] != tmp_contiguous_size * itemsize:
break
- tmp_contiguous_size *= a._shape[ndim-j-1]
+ tmp_contiguous_size *= out_shape[out_ndim-j-1]
contiguous_size = max(contiguous_size, tmp_contiguous_size)
return contiguous_size
@@ -368,7 +368,8 @@ cdef class _AbstractReductionKernel:
if optimize_context is None:
# Calculate manually
contiguous_size = _get_contiguous_size(
- in_args, self.in_params, in_shape.size(), out_shape.size())
+ in_args, self.in_params, out_shape,
+ in_shape.size(), out_shape.size())
block_size, block_stride, out_block_num = _get_block_specs(
internal.prod(in_shape),
internal.prod(out_shape),
@@ -407,7 +408,7 @@ cdef class _AbstractReductionKernel:
out_args = [_optimizer_copy_arg(a) for a in out_args]
contiguous_size = _get_contiguous_size(
- in_args, self.in_params, len(in_shape), len(out_shape))
+ in_args, self.in_params, out_shape, len(in_shape), len(out_shape))
block_size, block_stride, default_out_block_num = _get_block_specs(
internal.prod(in_shape),
internal.prod(out_shape),
