Repository: ivy-llc/ivy

Issue:
sparsify_tensor



Discussion:

Resolving patch:
diff --git a/ivy/data_classes/array/experimental/elementwise.py b/ivy/data_classes/array/experimental/elementwise.py
--- a/ivy/data_classes/array/experimental/elementwise.py
+++ b/ivy/data_classes/array/experimental/elementwise.py
@@ -1020,3 +1020,47 @@ def digamma(
         ivy.array([-0.7549271   0.92278427  0.9988394])
         """
         return ivy.digamma(self._data, out=out)
+
+    def sparsify_tensor(
+        self: ivy.Array,
+        card: int,
+        /,
+        *,
+        out: Optional[ivy.Array] = None,
+    ) -> ivy.Array:
+        """
+        ivy.Array class method variant of ivy.sparsify_tensor. This method simply wraps
+        the function, and so the docstring for ivy.sparsify_tensor also applies to this
+        method with minimal changes.
+
+        Parameters
+        ----------
+        self : array
+            The tensor to sparsify.
+        card : int
+            The number of values to keep.
+        out : array, optional
+            Optional output array, for writing the result to.
+
+        Returns
+        -------
+        ret : array
+            The sparsified tensor.
+
+        Examples
+        --------
+        >>> x = ivy.arange(100)
+        >>> x = ivy.reshape(x, (10, 10))
+        >>> x.sparsify_tensor(10)
+        ivy.array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+        [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])
+        """
+        return ivy.sparsify_tensor(self._data, card, out=out)
diff --git a/ivy/data_classes/container/experimental/elementwise.py b/ivy/data_classes/container/experimental/elementwise.py
--- a/ivy/data_classes/container/experimental/elementwise.py
+++ b/ivy/data_classes/container/experimental/elementwise.py
@@ -2975,3 +2975,96 @@ def digamma(
             map_sequences=map_sequences,
             out=out,
         )
+
+    @staticmethod
+    def static_sparsify_tensor(
+        x: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        card: Union[int, ivy.Container],
+        /,
+        *,
+        key_chains: Optional[Union[List[str], Dict[str, str]]] = None,
+        to_apply: Union[bool, ivy.Container] = True,
+        prune_unapplied: Union[bool, ivy.Container] = False,
+        map_sequences: Union[bool, ivy.Container] = False,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """
+        ivy.Container static method variant of ivy.sparsify_tensor. This method simply
+        wraps the function, and so the docstring for ivy.sparsify_tensor also applies to
+        this method with minimal changes.
+
+        Parameters
+        ----------
+        x
+            Input container containing input arrays.
+        card
+            The number of values to keep in each tensor.
+        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
+            Alternate output container in which to place the result.
+            The default is None.
+
+        Returns
+        -------
+        ret
+            container including the sparsified tensor computed element-wise
+        Examples
+        --------
+        >>> x = ivy.Container(
+                a=ivy.reshape(ivy.arange(100), (10, 10)),
+                b=ivy.reshape(ivy.arange(100), (10, 10)),
+            )
+        >>> ivy.Container.static_sparsify_tensor(x, 10)
+            {
+                a: (<class ivy.data_classes.array.array.Array> shape=[10, 10]),
+                b: (<class ivy.data_classes.array.array.Array> shape=[10, 10])
+            }
+        """
+        return ContainerBase.cont_multi_map_in_function(
+            "sparsify_tensor",
+            x,
+            card,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
+
+    def sparsify_tensor(
+        self: Union[ivy.Container, ivy.Array, ivy.NativeArray],
+        card: Union[int, ivy.Container],
+        /,
+        *,
+        key_chains: Optional[Union[List[str], Dict[str, str]]] = None,
+        to_apply: Union[bool, ivy.Container] = True,
+        prune_unapplied: Union[bool, ivy.Container] = False,
+        map_sequences: Union[bool, ivy.Container] = False,
+        out: Optional[ivy.Container] = None,
+    ) -> ivy.Container:
+        """
+        ivy.Container instance method variant of ivy.sparsify_tensor.
+
+        This method simply wraps the function, and so the docstring for
+        ivy.sparsify_tensor also applies to this method with minimal
+        changes.
+        """
+        return self.static_sparsify_tensor(
+            self,
+            card,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
diff --git a/ivy/functional/ivy/experimental/elementwise.py b/ivy/functional/ivy/experimental/elementwise.py
--- a/ivy/functional/ivy/experimental/elementwise.py
+++ b/ivy/functional/ivy/experimental/elementwise.py
@@ -1335,3 +1335,56 @@ def digamma(
     ivy.array([-0.7549271   0.92278427  0.9988394])
     """
     return ivy.current_backend(x).digamma(x, out=out)
+
+
+@handle_exceptions
+@handle_nestable
+@inputs_to_ivy_arrays
+@handle_array_function
+def sparsify_tensor(
+    x: Union[ivy.Array, ivy.NativeArray],
+    card: int,
+    /,
+    *,
+    out: Optional[ivy.Array] = None,
+) -> ivy.Array:
+    """
+    Zeros out all elements in the tensor except `card` elements with maximum absolute
+    values.
+
+    Parameters
+    ----------
+    x
+        Tensor to be sparsified
+    card
+        Desired number of non-zero elements in the tensor
+    out
+        Optional output array for writing the result to.
+
+    Returns
+    -------
+    ivy.array of shape tensor.shape
+
+    Examples
+    --------
+    >>> x = ivy.arange(100)
+    >>> x = ivy.reshape(x, (10, 10))
+    >>> sparsify_tensor(x, 10)
+    ivy.array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
+    [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])
+    """
+    if card >= ivy.prod(ivy.array(x.shape)):
+        return ivy.inplace_update(out, x) if ivy.exists(out) else x
+    _shape = ivy.shape(x)
+    x = ivy.reshape(ivy.sort(ivy.abs(x)), (-1,))
+    tensor = ivy.concat([ivy.zeros(len(x) - card, dtype=x.dtype), x[-card:]], axis=0)
+
+    return ivy.reshape(tensor, _shape, out=out)


Test patch:
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
--- 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
@@ -822,3 +822,49 @@ def test_digamma(
         on_device=on_device,
         x=x[0],
     )
+
+
+@st.composite
+def _sparsify_tensor_stg(draw):
+    dtype, tensor, shape = draw(
+        helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("numeric"),
+            ret_shape=True,
+            min_num_dims=1,
+            min_dim_size=1,
+            min_value=10,
+        )
+    )
+
+    size = 1
+    for dim in shape:
+        size *= dim
+
+    card = draw(st.integers(min_value=1, max_value=size))
+
+    return dtype, tensor[0], card
+
+
+# sparsify_tensor
+@handle_test(
+    fn_tree="functional.ivy.experimental.sparsify_tensor",
+    tensor_data=_sparsify_tensor_stg(),
+)
+def test_sparsify_tensor(
+    tensor_data,
+    test_flags,
+    on_device,
+    fn_name,
+    backend_fw,
+):
+    dtype, tensor, card = tensor_data
+
+    helpers.test_function(
+        backend_to_test=backend_fw,
+        test_flags=test_flags,
+        on_device=on_device,
+        fn_name=fn_name,
+        input_dtypes=dtype,
+        tensor=tensor,
+        card=card,
+    )