| from __future__ import annotations | |
| from cupy import _core | |
| 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) | |