refactor: added support_native_out attr & update.. (#10736)

Co-authored by MahmoudAshraf97 <[email protected]>

ivy/array/general.py

@@ -1199,7 +1199,8 @@ def scatter_flat(
         updates
             Values for the new array to hold.
         size
-            The size of the result.
+            The size of the result. Default is `None`, in which case tensor
+            argument out must be provided.
         reduction
             The reduction method for the scatter, one of 'sum', 'min', 'max' or
             'replace'
@@ -1211,6 +1212,25 @@ def scatter_flat(
         -------
         ret
             New array of given shape, with the values scattered at the indices.
+        
+        Examples
+        --------
+        With :class:`ivy.Array` input:
+        >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3])
+        >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3])
+        >>> size = 8
+        >>> out = indices.scatter_flat(updates, size=size)
+        >>> print(out)
+        ivy.array([8, 7, 5, 4, 0, 0, 0, 0])
+
+
+        With :class:`ivy.Array` input:
+        >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3])
+        >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3])
+        >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0])
+        >>> indices.scatter_flat(updates, out=out)
+        >>> print(out)
+        ivy.array([8, 7, 5, 4, 0, 0, 0, 0])
         """
         return ivy.scatter_flat(self, updates, size=size, reduction=reduction, out=out)
 

ivy/container/general.py

@@ -1852,7 +1852,8 @@ def static_scatter_flat(
         updates
             values to update input tensor with
         size
-            The size of the result.
+            The size of the result. Default is `None`, in which case tensor
+            argument out must be provided.
         reduction
             The reduction method for the scatter, one of 'sum', 'min', 'max'
             or 'replace'
@@ -1915,7 +1916,8 @@ def scatter_flat(
         updates
             values to update input tensor with
         size
-            The size of the result.
+            The size of the result. Default is `None`, in which case tensor
+            argument out must be provided.
         reduction
             The reduction method for the scatter, one of 'sum', 'min', 'max'
             or 'replace'
@@ -1938,6 +1940,20 @@ def scatter_flat(
         -------
         ret
             New container of given shape, with the values updated at the indices.
+
+        Examples
+        --------
+        With :class:`ivy.Container` input:
+        >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), \
+                            b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3]))
+        >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), \
+                            b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3]))
+        >>> size = 8
+        >>> print(ivy.scatter_flat(indices, updates, size=size))
+        {
+            a: ivy.array([4, 9, 5, 9, 0, 0, 0, 0]),
+            b: ivy.array([8, 7, 5, 4, 0, 0, 0, 0])
+        }
         """
         return self.static_scatter_flat(
             self,

ivy/functional/backends/jax/general.py

@@ -312,6 +312,9 @@ def scatter_flat(
     return _to_device(target)
 
 
+scatter_flat.support_native_out = True
+
+
 def scatter_nd(
     indices: JaxArray,
     updates: JaxArray,

ivy/functional/backends/numpy/general.py

@@ -266,6 +266,9 @@ def scatter_flat(
     return _to_device(target)
 
 
+scatter_flat.support_native_out = True
+
+
 def scatter_nd(
     indices: np.ndarray,
     updates: np.ndarray,

ivy/functional/backends/tensorflow/general.py

@@ -311,6 +311,9 @@ def scatter_flat(
     return res
 
 
+scatter_flat.support_native_out = True
+
+
 @with_unsupported_dtypes({"2.9.1 and below": ("bfloat16",)}, backend_version)
 def scatter_nd(
     indices: Union[tf.Tensor, tf.Variable],

ivy/functional/backends/torch/general.py

@@ -307,7 +307,7 @@ def scatter_flat(
     size: Optional[int] = None,
     reduction: str = "sum",
     out: Optional[torch.Tensor] = None,
-):
+) -> torch.Tensor:
     target = out
     target_given = ivy.exists(target)
     if ivy.exists(size) and ivy.exists(target):
@@ -354,6 +354,9 @@ def scatter_flat(
     return res
 
 
+scatter_flat.support_native_out = True
+
+
 @with_unsupported_dtypes(
     {
         "1.11.0 and below": (

ivy/functional/ivy/general.py

@@ -2775,7 +2775,8 @@ def scatter_flat(
     updates
         Values for the new array to hold.
     size
-        The size of the result.
+        The size of the result. Default is `None`, in which case tensor
+        argument out must be provided.
     reduction
         The reduction method for the scatter, one of 'sum', 'min', 'max' or 'replace'
     out
@@ -2787,6 +2788,53 @@ def scatter_flat(
     ret
         New array of given shape, with the values scattered at the indices.
 
+    This function is *nestable*, and therefore also accepts :code:'ivy.Container'
+    instance in place of the argument.
+
+    Examples
+    --------
+
+    With :class:`ivy.Array` input:
+    >>> indices = ivy.array([0, 0, 1, 0, 2, 2, 3, 3])
+    >>> updates = ivy.array([5, 1, 7, 2, 3, 2, 1, 3])
+    >>> out = ivy.array([0, 0, 0, 0, 0, 0, 0, 0])
+    >>> ivy.scatter_flat(indices, updates, out=out)
+    >>> print(out)
+    ivy.array([8, 7, 5, 4, 0, 0, 0, 0])
+
+
+    With :class:`ivy.Array` input:
+    >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3])
+    >>> updates = ivy.array([9, 2, 0, 2, 3, 2, 1, 8])
+    >>> size = 8
+    >>> print(ivy.scatter_flat(indices, updates, size=size))
+    ivy.array([4, 9, 5, 9, 0, 0, 0, 0])
+
+
+    With :class:`ivy.Container` and :class:`ivy.Array` input:
+    >>> indices = ivy.array([1, 0, 1, 0, 2, 2, 3, 3])
+    >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), \
+                        b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3]))
+    >>> size = 8
+    >>> print(ivy.scatter_flat(indices, updates, size=size))
+    {
+        a: ivy.array([4, 9, 5, 9, 0, 0, 0, 0]),
+        b: ivy.array([3, 12, 5, 4, 0, 0, 0, 0])
+    }
+
+
+    With :class:`ivy.Container` input:
+    >>> indices = ivy.Container(a=ivy.array([1, 0, 1, 0, 2, 2, 3, 3]), \
+                        b=ivy.array([0, 0, 1, 0, 2, 2, 3, 3]))
+    >>> updates = ivy.Container(a=ivy.array([9, 2, 0, 2, 3, 2, 1, 8]), \
+                        b=ivy.array([5, 1, 7, 2, 3, 2, 1, 3]))
+    >>> size = 8
+    >>> print(ivy.scatter_flat(indices, updates, size=size))
+    {
+        a: ivy.array([4, 9, 5, 9, 0, 0, 0, 0]),
+        b: ivy.array([8, 7, 5, 4, 0, 0, 0, 0])
+    }
+
     """
     return current_backend(indices).scatter_flat(
         indices, updates, size=size, reduction=reduction, out=out