support numpy array as output
@@ -525,6 +525,7 @@ class Plan1d(object):
CUFFT_COPY_HOST_TO_DEVICE)
check_result(result)
elif action == 'gather':
+ if isinstance(a, cupy.ndarray):
if self.batch == 1:
_reorder_buffers(plan, self.xtArr, xtArr_buffer)
@@ -535,6 +536,13 @@ class Plan1d(object):
xtArr_buffer[i], size)
start += count
assert start == b.size
+ else: # numpy
+ ptr2 = b.ctypes.data
+ with nogil:
+ result = cufftXtMemcpy(
+ plan, <void*>ptr2, <void*>arr,
+ CUFFT_COPY_DEVICE_TO_HOST)
+ check_result(result)
else:
raise ValueError
