def flip(a, axis): """Reverse the order of elements in an array along the given axis. Note that ``flip`` function has been introduced since NumPy v1.12. The contents of this document is the same as the original one. Args: a (~cupy.ndarray): Input array. axis (int): Axis in array, which entries are reversed. Returns: ~cupy.ndarray: Output array. .. seealso:: :func:`numpy.flip` """ a_ndim = a.ndim if a_ndim < 1: raise core._AxisError('Input must be >= 1-d') axis = int(axis) if not -a_ndim <= axis < a_ndim: raise core._AxisError('axis must be >= %d and < %d' % (-a_ndim, a_ndim)) return _flip(a, axis)
def _cum_core(a, axis, dtype, out, kern, batch_kern): a = cupy.asarray(a) if out is None: if dtype is None: kind = a.dtype.kind if kind == 'b': dtype = numpy.dtype('l') elif kind == 'i' and a.dtype.itemsize < numpy.dtype('l').itemsize: dtype = numpy.dtype('l') elif kind == 'u' and a.dtype.itemsize < numpy.dtype('L').itemsize: dtype = numpy.dtype('L') else: dtype = a.dtype out = a.astype(dtype) else: out[...] = a if axis is None: out = out.ravel() elif not (-a.ndim <= axis < a.ndim): raise core._AxisError('axis(={}) out of bounds'.format(axis)) else: return _proc_as_batch(batch_kern, out, axis=axis) pos = 1 while pos < out.size: kern(pos, out, size=out.size) pos <<= 1 return out
def _get_positive_axis(ndim, axis): a = axis if a < 0: a += ndim if a < 0 or a >= ndim: raise core._AxisError( 'axis {} out of bounds [0, {})'.format(axis, ndim)) return a
def stack(tup, axis=0): """Stacks arrays along a new axis. Args: tup (sequence of arrays): Arrays to be stacked. axis (int): Axis along which the arrays are stacked. Returns: cupy.ndarray: Stacked array. .. seealso:: :func:`numpy.stack` """ # TODO(okuta) Remove this if exampd_dims is updated for x in tup: if not (-x.ndim - 1 <= axis <= x.ndim): raise core._AxisError('axis {} out of bounds [{}, {}]'.format( axis, -x.ndim - 1, x.ndim)) return concatenate([cupy.expand_dims(x, axis) for x in tup], axis)
def _cum_core(a, axis, dtype, out, kern, batch_kern, *, op=None): a = cupy.asarray(a) if out is None: if dtype is None: kind = a.dtype.kind if kind == 'b': dtype = numpy.dtype('l') elif kind == 'i' and a.dtype.itemsize < numpy.dtype('l').itemsize: dtype = numpy.dtype('l') elif kind == 'u' and a.dtype.itemsize < numpy.dtype('L').itemsize: dtype = numpy.dtype('L') else: dtype = a.dtype out = a.astype(dtype) else: out[...] = a if axis is None: out = out.ravel() elif not (-a.ndim <= axis < a.ndim): raise core._AxisError('axis(={}) out of bounds'.format(axis)) else: return _proc_as_batch(batch_kern, out, axis=axis) if cupy.cuda.cub_enabled: # result will be None if the scan is not compatible with CUB result = cub.cub_scan(out, op) if result is not None: return result pos = 1 while pos < out.size: kern(pos, out, size=out.size) pos <<= 1 return out