diff --git a/ivy/array/experimental/elementwise.py b/ivy/array/experimental/elementwise.py index 6dca0389761c6..2a4dbde71dfab 100644 --- a/ivy/array/experimental/elementwise.py +++ b/ivy/array/experimental/elementwise.py @@ -1,6 +1,6 @@ # global import abc -from typing import Optional, Union +from typing import Optional, Union, Tuple # local import ivy @@ -274,6 +274,59 @@ def exp2( """ return ivy.exp2(self._data, out=out) + def count_nonzero( + self: ivy.Array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[ivy.Array] = None, + ) -> ivy.Array: + """ivy.Array instance method variant of ivy.count_nonzero. This method simply + wraps the function, and so the docstring for ivy.count_nonzero also applies to + this method with minimal changes. + + Parameters + ---------- + self + input array for which to count non-zeros. + axis + 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 the input array. + keepdims + optional, if this is set to True, the axes that are counted are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + dtype + optional output dtype. Default is of type integer. + out + optional output array, for writing the result to. + + Returns + ------- + ret + 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. + + Examples + -------- + >>> x = ivy.array([1, 2, 3]) + >>> x.count_nonzero() + ivy.array(3) + >>> x = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x.count_nonzero(axis=0) + ivy.array([[1, 2], + [2, 2]]) + >>> x = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> x.count_nonzero(axis=(0,1), keepdims=True) + ivy.array([[[3, 4]]]) + """ + return ivy.count_nonzero( + self._data, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) + def nansum( self: ivy.Array, /, diff --git a/ivy/container/experimental/elementwise.py b/ivy/container/experimental/elementwise.py index c630f03d0c825..f09317f286d36 100644 --- a/ivy/container/experimental/elementwise.py +++ b/ivy/container/experimental/elementwise.py @@ -1,5 +1,5 @@ # global -from typing import Optional, Union, List, Dict +from typing import Optional, Union, List, Dict, Tuple # local import ivy @@ -604,6 +604,187 @@ def exp2( """ return self.static_exp2(self, out=out) + @staticmethod + def static_count_nonzero( + a: Union[ivy.Array, ivy.NativeArray, ivy.Container], + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + key_chains: Optional[Union[List[str], Dict[str, str]]] = None, + to_apply: bool = True, + prune_unapplied: bool = False, + map_sequences: bool = False, + out: Optional[ivy.Container] = None, + ) -> ivy.Container: + """ + ivy.Container static method variant of ivy.count_nonzero. This method simply + wraps the function, and so the docstring for ivy.count_nonzero also applies + to this method with minimal changes. + + Parameters + ---------- + a + container with the base input arrays. + axis + 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 the input array. + keepdims + optional, if this is set to True, the axes that are counted are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + dtype + optional output dtype. Default is of type integer. + key_chains + The key-chains to apply or not apply the method to. Default is None. + to_apply + If True, the method will be applied to key_chains, otherwise key_chains + will be skipped. Default is True. + prune_unapplied + Whether to prune key_chains for which the function was not applied. + Default is False. + map_sequences + Whether to also map method to sequences (lists, tuples). Default is False. + out + optional output container, for writing the result to. + + Returns + ------- + ret + Container including number of non-zero values in the array along a + given axis. Otherwise, container with the total number of non-zero + values in the array is returned. + + Examples + -------- + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> ivy.Container.static_count_nonzero(x) + { + a: ivy.array(7), + b: ivy.array(7) + } + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> ivy.Container.static_count_nonzero(x, axis=0) + { + a: ivy.array([1, 2, 2, 2]), + b: ivy.array([[1, 2], + [2, 2]]) + } + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> ivy.Container.static_count_nonzero(x, axis=(0,1), keepdims=True) + { + a: ivy.array([[7]]), + b: ivy.array([[[3, 4]]]) + } + """ + return ContainerBase.multi_map_in_static_method( + "count_nonzero", + a, + axis=axis, + keepdims=keepdims, + dtype=dtype, + key_chains=key_chains, + to_apply=to_apply, + prune_unapplied=prune_unapplied, + map_sequences=map_sequences, + out=out, + ) + + def count_nonzero( + self: ivy.Container, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + key_chains: Optional[Union[List[str], Dict[str, str]]] = None, + to_apply: bool = True, + prune_unapplied: bool = False, + map_sequences: bool = False, + out: Optional[ivy.Container] = None, + ) -> ivy.Container: + """ + ivy.Container instance method variant of ivy.count_nonzero. This method + simply wraps the function, and so the docstring for ivy.count_nonzero also + applies to this method with minimal changes. + + Parameters + ---------- + self + container with the base input arrays. + axis + 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 the input array. + keepdims + optional, if this is set to True, the axes that are counted are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + dtype + optional output dtype. Default is of type integer. + key_chains + The key-chains to apply or not apply the method to. Default is ``None``. + to_apply + If True, the method will be applied to key_chains, otherwise key_chains + will be skipped. Default is ``True``. + prune_unapplied + Whether to prune key_chains for which the function was not applied. + Default is ``False``. + map_sequences + Whether to also map method to sequences (lists, tuples). + Default is ``False`` + out + optional output container, for writing the result to. + + Returns + ------- + ret + Container including number of non-zero values in the array along a + given axis. Otherwise, container with the total number of non-zero + values in the array is returned. + + Examples + -------- + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> x.count_nonzero() + { + a: ivy.array(7), + b: ivy.array(7) + } + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> x.count_nonzero(axis=0) + { + a: ivy.array([1, 2, 2, 2]), + b: ivy.array([[1, 2], + [2, 2]]) + } + >>> x = ivy.Container(a=ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]),\ + b=ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]])) + >>> x.count_nonzero(axis=(0,1), keepdims=True) + { + a: ivy.array([[7]]), + b: ivy.array([[[3, 4]]]) + } + """ + return self.static_count_nonzero( + self, + axis=axis, + keepdims=keepdims, + dtype=dtype, + key_chains=key_chains, + to_apply=to_apply, + prune_unapplied=prune_unapplied, + map_sequences=map_sequences, + out=out, + ) + @staticmethod def static_nansum( x: Union[ivy.Container, ivy.Array, ivy.NativeArray], diff --git a/ivy/functional/backends/jax/experimental/elementwise.py b/ivy/functional/backends/jax/experimental/elementwise.py index 668bbbd27bba0..505df90a97b6e 100644 --- a/ivy/functional/backends/jax/experimental/elementwise.py +++ b/ivy/functional/backends/jax/experimental/elementwise.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Optional, Union, Tuple from ivy.functional.backends.jax import JaxArray import jax.numpy as jnp @@ -62,6 +62,22 @@ def exp2( return jnp.exp2(x) +def count_nonzero( + a: JaxArray, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[jnp.dtype] = None, + out: Optional[JaxArray] = None, +) -> JaxArray: + if isinstance(axis, list): + axis = tuple(axis) + if dtype is None: + return jnp.count_nonzero(a, axis=axis, keepdims=keepdims) + return jnp.array(jnp.count_nonzero(a, axis=axis, keepdims=keepdims), dtype=dtype) + + def nansum( x: JaxArray, /, diff --git a/ivy/functional/backends/numpy/experimental/elementwise.py b/ivy/functional/backends/numpy/experimental/elementwise.py index 2f9893d611eec..7e6034be428ea 100644 --- a/ivy/functional/backends/numpy/experimental/elementwise.py +++ b/ivy/functional/backends/numpy/experimental/elementwise.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Optional, Union, Tuple import numpy as np from ivy.functional.backends.numpy.helpers import _scalar_output_to_0d_array @@ -109,6 +109,26 @@ def exp2( exp2.support_native_out = True +@_scalar_output_to_0d_array +def count_nonzero( + x: np.ndarray, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[np.dtype] = None, + out: Optional[np.ndarray] = None, +) -> np.ndarray: + if isinstance(axis, list): + axis = tuple(axis) + if dtype is None: + return np.count_nonzero(x, axis=axis, keepdims=keepdims) + return np.array(np.count_nonzero(x, axis=axis, keepdims=keepdims), dtype=dtype) + + +count_nonzero.support_native_out = False + + def nansum( x: np.ndarray, /, diff --git a/ivy/functional/backends/tensorflow/experimental/elementwise.py b/ivy/functional/backends/tensorflow/experimental/elementwise.py index 0ab67078fb894..796a93eaf111a 100644 --- a/ivy/functional/backends/tensorflow/experimental/elementwise.py +++ b/ivy/functional/backends/tensorflow/experimental/elementwise.py @@ -1,4 +1,4 @@ -from typing import Union, Optional +from typing import Union, Optional, Tuple import tensorflow as tf from .. import backend_version @@ -95,6 +95,22 @@ def exp2( return tf.math.pow(2, x, name=None) +def count_nonzero( + a: Union[tf.Tensor, tf.Variable], + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[tf.DType] = None, + out: Optional[Union[tf.Tensor, tf.Variable]] = None, +) -> Union[tf.Tensor, tf.Variable]: + if dtype is None: + return tf.math.count_nonzero(a, axis=axis, keepdims=keepdims, name=None) + return tf.math.count_nonzero( + a, axis=axis, keepdims=keepdims, dtype=dtype, name=None + ) + + def nansum( x: Union[tf.Tensor, tf.Variable], /, diff --git a/ivy/functional/backends/torch/experimental/elementwise.py b/ivy/functional/backends/torch/experimental/elementwise.py index c555e9ee0de7f..3c65a57578381 100644 --- a/ivy/functional/backends/torch/experimental/elementwise.py +++ b/ivy/functional/backends/torch/experimental/elementwise.py @@ -1,5 +1,5 @@ # global -from typing import Optional, Union +from typing import Optional, Union, Tuple import torch # local @@ -109,6 +109,36 @@ def exp2( exp2.support_native_out = True +def count_nonzero( + a: torch.Tensor, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[torch.dtype] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if isinstance(axis, list): + axis = tuple(axis) + + def _dtype_count_nonzero(a, axis, dtype): + if dtype is None: + return torch.count_nonzero(a, dim=axis) + return torch.tensor(torch.count_nonzero(a, dim=axis), dtype=dtype) + + x = _dtype_count_nonzero(a, axis, dtype) + if not keepdims: + return x + if isinstance(axis, tuple): + for d in sorted(axis, reverse=True): + x = x.unsqueeze(d) + return x + return x.unsqueeze(axis) + + +count_nonzero.support_native_out = False + + def nansum( x: torch.Tensor, /, diff --git a/ivy/functional/experimental/elementwise.py b/ivy/functional/experimental/elementwise.py index d8fc2f52c9862..56051bdf599c8 100644 --- a/ivy/functional/experimental/elementwise.py +++ b/ivy/functional/experimental/elementwise.py @@ -1,5 +1,5 @@ # local -from typing import Optional, Union +from typing import Optional, Union, Tuple import ivy from ivy.func_wrapper import ( handle_out_argument, @@ -343,6 +343,61 @@ def exp2( return ivy.current_backend().exp2(x, out=out) +@to_native_arrays_and_back +@handle_out_argument +@handle_nestable +@handle_exceptions +def count_nonzero( + a: Union[ivy.Array, ivy.NativeArray], + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + keepdims: Optional[bool] = False, + dtype: Optional[Union[ivy.Dtype, ivy.NativeDtype]] = None, + out: Optional[Union[int, ivy.Array]] = None, +) -> ivy.Array: + """Counts the number of non-zero values in the array a. + + Parameters + ---------- + a + array for which to count non-zeros. + axis + 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 the input array. + keepdims + optional, if this is set to True, the axes that are counted are left in the + result as dimensions with size one. With this option, the result + will broadcast correctly against the input array. + dtype + optional output dtype. Default is of type integer. + out + optional output array, for writing the result to. + + Returns + ------- + ret + 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. + + Examples + -------- + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a) + ivy.array(7) + >>> a = ivy.array([[0, 1, 2, 3],[4, 5, 6, 7]]) + >>> ivy.count_nonzero(a, axis=0) + ivy.array([1, 2, 2, 2]) + >>> a = ivy.array([[[0,1],[2,3]],[[4,5],[6,7]]]) + >>> ivy.count_nonzero(a, axis=(0,1), keepdims=True) + ivy.array([[[3, 4]]]) + """ + return ivy.current_backend().count_nonzero( + a, axis=axis, keepdims=keepdims, dtype=dtype, out=out + ) + + @to_native_arrays_and_back @handle_out_argument @handle_nestable diff --git a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_elementwise.py b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_elementwise.py index 62af045a362e6..6bfb10c69a8cb 100644 --- a/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_elementwise.py +++ b/ivy_tests/test_ivy/test_functional/test_experimental/test_core/test_elementwise.py @@ -338,6 +338,74 @@ def test_exp2( ) +@st.composite +def _get_dtype_values_axis_for_count_nonzero( + draw, + in_available_dtypes, + out_available_dtypes, + min_num_dims=1, + max_num_dims=10, + min_dim_size=1, + max_dim_size=10, +): + input_dtype, values, axis = draw( + helpers.dtype_values_axis( + available_dtypes=helpers.get_dtypes(in_available_dtypes), + min_num_dims=min_num_dims, + max_num_dims=max_num_dims, + min_dim_size=min_dim_size, + max_dim_size=max_dim_size, + valid_axis=True, + ) + ) + axis = draw(st.one_of(st.just(axis), st.none())) + output_dtype = draw(helpers.get_dtypes(out_available_dtypes)) + return [input_dtype, output_dtype], values, axis + + +# count_nonzero +@handle_cmd_line_args +@given( + dtype_values_axis=_get_dtype_values_axis_for_count_nonzero( + in_available_dtypes="numeric", + out_available_dtypes="numeric", + min_num_dims=1, + max_num_dims=10, + min_dim_size=1, + max_dim_size=10, + ), + keepdims=st.booleans(), + num_positional_args=helpers.num_positional_args(fn_name="count_nonzero"), +) +def test_count_nonzero( + dtype_values_axis, + keepdims, + with_out, + as_variable, + num_positional_args, + native_array, + container, + instance_method, + fw, +): + i_o_dtype, a, axis = dtype_values_axis + helpers.test_function( + input_dtypes=i_o_dtype[0], + as_variable_flags=as_variable, + with_out=with_out, + num_positional_args=num_positional_args, + native_array_flags=native_array, + container_flags=container, + instance_method=instance_method, + fw=fw, + fn_name="count_nonzero", + a=a, + axis=axis, + keepdims=keepdims, + dtype=i_o_dtype[1], + ) + + # nansum @handle_cmd_line_args @given(