This RFC proposes a new addition to the array API specification for counting the number of "non-zero" (i.e., truthy) values in an array.

Overview

Based on array comparison data, the API is available across all major array libraries in the PyData ecosystem.

count_nonzero was originally identified in #187 as a potential standardization candidate and has usage within downstream libraries (e.g., sklearn, SciPy).

Prior art

• NumPy: https://numpy.org/doc/stable/reference/generated/numpy.count_nonzero.html
• CuPy: https://docs.cupy.dev/en/stable/reference/generated/cupy.count_nonzero.html
  • does not support keepdims.
• Dask: https://docs.dask.org/en/stable/generated/dask.array.count_nonzero.html
  • does not support keepdims.
• JAX: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.count_nonzero.html
• PyTorch: https://pytorch.org/docs/stable/generated/torch.count_nonzero.html
  • does not support keepdims and was not discussed upon initial addition (ref).
• TensorFlow: https://www.tensorflow.org/api_docs/python/tf/math/count_nonzero

Proposal

def count_nonzero(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array

• When axis is None, the function should count the number of non-zero elements along a flattened array.
• The function should return an array having the default index data type.

Questions

• In contrast to sum and other reductions, support for keepdims is less common among array libraries. Why this is the case is not clear. Are there any reasons why keepdims should not be standardized?