add 1d kernel
@@ -74,6 +74,11 @@ cpdef ndarray _ndarray_argwhere(ndarray self):
if count_nonzero == 0:
return ndarray((0, ndim), dtype=dtype)
+ if ndim <= 1:
+ dst = ndarray((count_nonzero, 1), dtype=dtype)
+ _nonzero_kernel_1d(nonzero, scan_index, dst)
+ return dst
+ else:
nonzero.shape = self.shape
scan_index.shape = self.shape
dst = ndarray((count_nonzero, ndim), dtype=dtype)
@@ -355,6 +360,12 @@ cdef ndarray _simple_getitem(ndarray a, list slice_list):
return v
+_nonzero_kernel_1d = ElementwiseKernel(
+ 'T src, S index', 'raw S dst',
+ 'if (src != 0) dst[index - 1] = i',
+ 'nonzero_kernel_1d')
+
+
_nonzero_kernel = ElementwiseKernel(
'T src, S index', 'raw S dst',
'''
