def setUp(self):
if self.stream == 'null':
self.stream = cupy.cuda.Stream.null
elif self.stream == 'new':
self.stream = cupy.cuda.Stream()
self.my_int8_sum = _core.create_reduction_func(
'my_sum', ('b->b', ), ('in0', 'a + b', 'out0 = a', None))
def get_sum_func(self):
return _core.create_reduction_func('my_sum', ('b->b', ),
('in0', 'a + b', 'out0 = a', None),
0)
from cupy import _core
from cupy.linalg import _decomposition
from cupy.linalg import _util
import functools
def _multi_svd_norm(x, row_axis, col_axis, op):
y = cupy.moveaxis(x, (row_axis, col_axis), (-2, -1))
result = op(_decomposition.svd(y, compute_uv=False), axis=-1)
return result
_norm_ord2 = _core.create_reduction_func(
'_norm_ord2',
('?->l', 'b->l', 'B->L', 'h->l', 'H->L', 'i->l', 'I->L', 'l->l', 'L->L',
'q->q', 'Q->Q', ('e->e', (None, None, None, 'float')), 'f->f', 'd->d'),
('in0 * in0', 'a + b', 'out0 = sqrt(type_out0_raw(a))', None), 0)
_norm_ord2_complex = _core.create_reduction_func(
'_norm_ord2_complex', ('F->f', 'D->d'),
('in0.real() * in0.real() + in0.imag() * in0.imag()', 'a + b',
'out0 = sqrt(type_out0_raw(a))', None), 0)
def norm(x, ord=None, axis=None, keepdims=False):
"""Returns one of matrix norms specified by ``ord`` parameter.
See numpy.linalg.norm for more detail.
Args:
x (cupy.ndarray): Array to take norm. If ``axis`` is None,
def count_nonzero(a, axis=None):
"""Counts the number of non-zero values in the array.
.. note::
:func:`numpy.count_nonzero` returns `int` value when `axis=None`,
but :func:`cupy.count_nonzero` returns zero-dimensional array to reduce
CPU-GPU synchronization.
Args:
a (cupy.ndarray): The array for which to count non-zeros.
axis (int or tuple, optional): Axis or tuple of axes along which to
count non-zeros. Default is None, meaning that non-zeros will be
counted along a flattened version of ``a``
Returns:
cupy.ndarray of int: Number of non-zero values in the array
along a given axis. Otherwise, the total number of non-zero values
in the array is returned.
"""
return _count_nonzero(a, axis=axis)
_count_nonzero = _core.create_reduction_func(
'cupy_count_nonzero',
('?->l', 'B->l', 'h->l', 'H->l', 'i->l', 'I->l', 'l->l', 'L->l', 'q->l',
'Q->l', 'e->l', 'f->l', 'd->l', 'F->l', 'D->l'),
('in0 != type_in0_raw(0)', 'a + b', 'out0 = a', None), 0)