MAINT: count_nonzero

Author: ev-br

torch_np/_detail/_reductions.py

@@ -65,6 +65,8 @@ def _atleast_float(dtype, other_dtype):
     return dtype
 
 
+@emulate_keepdims
+@deco_axis_expand
 def count_nonzero(a, axis=None):
     # XXX: this all should probably be generalized to a sum(a != 0, dtype=bool)
     try:

torch_np/_funcs.py

@@ -442,3 +442,9 @@ def any(a: ArrayLike, axis: AxisLike=None, out=None, keepdims=NoValue, *, where=
     result = _reductions.any(a, axis=axis, where=where, keepdims=keepdims)
     return _helpers.result_or_out(result, out)
 
+
+@normalizer
+def count_nonzero(a: ArrayLike, axis: AxisLike=None, *, keepdims=False):
+    result = _reductions.count_nonzero(a, axis=axis, keepdims=keepdims)
+    return _helpers.array_from(result)
+

torch_np/_wrapper.py

@@ -582,12 +582,6 @@ def flatnonzero(a):
     return _funcs.nonzero(arr.ravel())[0]
 
 
-from ._decorators import emulate_out_arg
-from ._ndarray import axis_keepdims_wrapper
-
-count_nonzero = emulate_out_arg(axis_keepdims_wrapper(_reductions.count_nonzero))
-
-
 @normalizer
 def roll(a: ArrayLike, shift, axis=None):
     result = _impl.roll(a, shift, axis)