hstack extension (#6069)

* hstack initial setup

* examples for all docstrings

* lint fixes

ivy/array/extensions/manipulation.py

@@ -144,3 +144,28 @@ def vstack(
                        [7, 8]])
         """
         return ivy.vstack(self.concat(arrays), out=out)
+
+    def hstack(
+        self: ivy.Array,
+        /,
+        arrays: Union[
+            Tuple[Union[ivy.Array, ivy.NativeArray]],
+            List[Union[ivy.Array, ivy.NativeArray]],
+        ],
+        *,
+        out: Optional[ivy.Array] = None,
+    ) -> ivy.Array:
+        """
+        ivy.Array instance method variant of ivy.hstack. This method simply
+        wraps the function, and so the docstring for ivy.hstack also applies
+        to this method with minimal changes.
+
+        Examples
+        --------
+        >>> x = ivy.array([[1, 2]])
+        >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])]
+        >>> print(x.vstack(y))
+        ivy.array([1, 2, 5, 6, 7, 8])
+
+        """
+        return ivy.hstack(self.concat(arrays), out=out)

ivy/container/extensions/manipulation.py

@@ -328,15 +328,6 @@ def vstack(
             b: ivy.array([[[4, 5]],
                         [[1, 0]]])
         }
-        >>> ivy.Container.static_stack([x,y],axis=1)
-        {
-            a: ivy.array([[[0, 1],
-                        [3, 2]],
-                        [[2, 3],
-                        [1, 0]]]),
-            b: ivy.array([[[4, 5],
-                        [1, 0]]])
-        }
         """
         new_xs = xs.copy()
         new_xs.insert(0, self.copy())
@@ -391,3 +382,82 @@ def static_vstack(
             map_sequences=map_sequences,
             out=out,
         )
+
+    def hstack(
+        self: ivy.Container,
+        /,
+        xs: Union[
+            Tuple[Union[ivy.Array, ivy.NativeArray, ivy.Container]],
+            List[Union[ivy.Array, ivy.NativeArray, ivy.Container]],
+        ],
+        *,
+        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.hstack. This method
+        simply wraps the function, and so the docstring for ivy.hstack
+        also applies to this method with minimal changes.
+
+        Examples
+        --------
+        >>> x = ivy.Container(a=ivy.array([[0, 1], [2,3]]), b=ivy.array([[4, 5]]))
+        >>> y = ivy.Container(a=ivy.array([[3, 2], [1,0]]), b=ivy.array([[1, 0]]))
+        >>> x.hstack([y])
+        {
+            a: ivy.array([[0, 1, 3, 2],
+                          [2, 3, 1, 0]]),
+            b: ivy.array([[4, 5, 1, 0]])
+        }
+        """
+        new_xs = xs.copy()
+        new_xs.insert(0, self.copy())
+        return self.static_hstack(
+            new_xs,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )
+
+    @staticmethod
+    def static_hstack(
+        xs: Union[
+            Tuple[Union[ivy.Array, ivy.NativeArray, ivy.Container]],
+            List[Union[ivy.Array, ivy.NativeArray, ivy.Container]],
+        ],
+        /,
+        *,
+        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.hstack. This method simply wraps the
+        function, and so the docstring for ivy.hstack also applies to this method
+        with minimal changes.
+
+        Examples
+        --------
+        With one :class:`ivy.Container` input:
+        >>> c = ivy.Container(a=[ivy.array([1,2,3]), ivy.array([0,0,0])])
+        >>> ivy.Container.static_hstack(c)
+        {
+            a: ivy.array([1, 2, 3, 0, 0, 0])
+        }
+        """
+        return ContainerBase.multi_map_in_static_method(
+            "hstack",
+            xs,
+            key_chains=key_chains,
+            to_apply=to_apply,
+            prune_unapplied=prune_unapplied,
+            map_sequences=map_sequences,
+            out=out,
+        )

ivy/functional/backends/jax/extensions/manipulation.py

@@ -40,3 +40,12 @@ def vstack(
     out: Optional[JaxArray] = None,
 ) -> JaxArray:
     return jnp.vstack(arrays)
+
+
+def hstack(
+    arrays: Sequence[JaxArray],
+    /,
+    *,
+    out: Optional[JaxArray] = None,
+) -> JaxArray:
+    return jnp.hstack(arrays)

ivy/functional/backends/numpy/extensions/manipulation.py

@@ -52,3 +52,12 @@ def vstack(
     out: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     return np.vstack(arrays)
+
+
+def hstack(
+    arrays: Sequence[np.ndarray],
+    /,
+    *,
+    out: Optional[np.ndarray] = None,
+) -> np.ndarray:
+    return np.hstack(arrays)

ivy/functional/backends/tensorflow/extensions/manipulation.py

@@ -42,3 +42,12 @@ def vstack(
     out: Optional[Union[tf.Tensor, tf.Variable]] = None,
 ) -> Union[tf.Tensor, tf.Variable]:
     return tf.experimental.numpy.vstack(arrays)
+
+
+def hstack(
+    arrays: Union[Sequence[tf.Tensor], Sequence[tf.Variable]],
+    /,
+    *,
+    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
+) -> Union[tf.Tensor, tf.Variable]:
+    return tf.experimental.numpy.hstack(arrays)

ivy/functional/backends/torch/extensions/manipulation.py

@@ -52,3 +52,12 @@ def vstack(
     out: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     return torch.vstack(arrays, out=None)
+
+
+def hstack(
+    arrays: Sequence[torch.Tensor],
+    /,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    return torch.hstack(arrays, out=None)

ivy/functional/extensions/manipulation.py

@@ -368,7 +368,7 @@ def vstack(arrays: Sequence[ivy.Array], /) -> ivy.Array:
     ivy.array([[1, 2, 3],
            [2, 3, 4]])
     >>> ivy.vstack((x, y, x, y))
-        ivy.array([[1, 2, 3],
+    ivy.array([[1, 2, 3],
                [2, 3, 4],
                [1, 2, 3],
                [2, 3, 4]])
@@ -380,3 +380,37 @@ def vstack(arrays: Sequence[ivy.Array], /) -> ivy.Array:
 
     """
     return ivy.current_backend(arrays[0]).vstack(arrays)
+
+
+@to_native_arrays_and_back
+@handle_out_argument
+@handle_nestable
+def hstack(arrays: Sequence[ivy.Array], /) -> ivy.Array:
+    """Stack arrays in sequence horizotally (column wise).
+
+    Parameters
+    ----------
+    arrays
+        Sequence of arrays to be stacked.
+
+    Returns
+    -------
+    ret
+        The array formed by stacking the given arrays.
+
+    Examples
+    --------
+    >>> x = ivy.array([1, 2, 3])
+    >>> y = ivy.array([2, 3, 4])
+    >>> ivy.hstack((x, y))
+    ivy.array([1, 2, 3, 2, 3, 4])
+    >>> x = ivy.array([1, 2, 3])
+    >>> y = ivy.array([0, 0, 0])
+    >>> ivy.hstack((x, y, x))
+    ivy.array([1, 2, 3, 0, 0, 0, 1, 2, 3])
+    >>> y = [ivy.array([[5, 6]]), ivy.array([[7, 8]])]
+    >>> print(ivy.hstack(y))
+    ivy.array([[5, 6, 7, 8]])
+
+    """
+    return ivy.current_backend(arrays[0]).hstack(arrays)

ivy_tests/test_ivy/test_functional/test_extensions/test_core/test_manipulation.py

@@ -239,3 +239,44 @@ def test_vstack(
         fn_name="vstack",
         arrays=m,
     )
+
+
+@handle_cmd_line_args
+@given(
+    dtype_and_m=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("float"),
+        min_value=-10,
+        max_value=10,
+        shared_dtype=True,
+        num_arrays=2,
+        shape=helpers.get_shape(
+            min_num_dims=1,
+            max_num_dims=3,
+            min_dim_size=1,
+            max_dim_size=3),
+    ),
+    num_positional_args=helpers.num_positional_args(fn_name="hstack"),
+)
+def test_hstack(
+    dtype_and_m,
+    as_variable,
+    with_out,
+    num_positional_args,
+    native_array,
+    container,
+    instance_method,
+    fw,
+):
+    input_dtype, m = dtype_and_m
+    helpers.test_function(
+        input_dtypes=input_dtype,
+        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="hstack",
+        arrays=m,
+    )