Use out argument
@@ -621,16 +621,19 @@ cdef class ndarray:
choices = choices[:, None, ...]
ba, bcs = broadcast(a, choices).values
+ if out is None:
+ out = ndarray(ba.shape[1:], choices.dtype)
+
n_channel = numpy.prod(bcs[0].shape)
if mode == 'raise':
if not ((a < n).all() and (0 <= a).all()):
raise ValueError('invalid entry in choice array')
- c = _choose_kernel(ba[0], bcs, n_channel)
+ c = _choose_kernel(ba[0], bcs, n_channel, out)
elif mode == 'wrap':
ba = ba[0] % n
- c = _choose_kernel(ba, bcs, n_channel)
+ c = _choose_kernel(ba, bcs, n_channel, out)
elif mode == 'clip':
- c = _choose_clip_kernel(ba[0], bcs, n_channel, n)
+ c = _choose_clip_kernel(ba[0], bcs, n_channel, n, out)
else:
raise TypeError('clipmode not understood')
