[PT FE] Added support for aten::hstack and aten::vstack (#28933)

### Details:
 - *added aten::hstack*
 - *added aten::vstack*

### Tickets:
 - Closes #28875

---------

Co-authored-by: Maxim Vafin <[email protected]>
Co-authored-by: Przemyslaw Wysocki <[email protected]>

Author: 3 people

src/frontends/pytorch/src/op/cat.cpp

@@ -155,6 +155,28 @@ OutputVector translate_stack_fx(const NodeContext& context) {
     return translate_cat_common(context, list_elems, axis, true);
 }
 
+OutputVector translate_hstack(const NodeContext& context) {
+    num_inputs_check(context, 1, 2);
+    const auto&& list_elems = get_list_as_outputs(context.get_input(0));
+    int64_t axis = 1;
+    auto out = translate_cat_common(context, list_elems, axis, false);
+    if (!context.input_is_none(1)) {
+        context.mutate_input(1, out[0]);
+    }
+    return out;
+};
+
+OutputVector translate_vstack(const NodeContext& context) {
+    num_inputs_check(context, 1, 2);
+    const auto&& list_elems = get_list_as_outputs(context.get_input(0));
+    int64_t axis = 0;
+    auto out = translate_cat_common(context, list_elems, axis, false);
+    if (!context.input_is_none(1)) {
+        context.mutate_input(1, out[0]);
+    }
+    return out;
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/op_table.cpp

@@ -124,6 +124,7 @@ OP_CONVERTER(translate_group_norm);
 OP_CONVERTER(translate_gru);
 OP_CONVERTER(translate_hann_window);
 OP_CONVERTER(translate_hardtanh);
+OP_CONVERTER(translate_hstack);
 OP_CONVERTER(translate_if);
 OP_CONVERTER(translate_im2col);
 OP_CONVERTER(translate_index);
@@ -341,6 +342,7 @@ OP_CONVERTER(translate_quantize_per_channel_fx);
 OP_CONVERTER(translate_quantize_per_tensor_fx);
 OP_CONVERTER(translate_var_fx);
 OP_CONVERTER(translate_var_mean_fx);
+OP_CONVERTER(translate_vstack);
 OP_CONVERTER(translate_unbind_int_fx);
 OP_CONVERTER(translate_zeros_fx);
 OP_CONVERTER(translate_zeros_like_fx);
@@ -531,6 +533,7 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::hardsigmoid", op::quantizable_op<op::translate_1to1_match_1_inputs<opset10::HSigmoid>>},
         {"aten::hardswish", op::quantizable_op<op::translate_1to1_match_1_inputs<opset10::HSwish>>},
         {"aten::hardtanh", op::quantizable_op<op::translate_hardtanh>},
+        {"aten::hstack", op::translate_hstack},
         {"aten::im2col", op::translate_im2col},
         {"aten::imag", common_translators::translate_imag},
         // aten::index - Supported in limited set of patterns
@@ -737,6 +740,7 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::view_as", op::translate_reshape_as},
         {"aten::view_as_complex", op::translate_view_as_complex},
         {"aten::view_as_real", op::translate_view_as_real},
+        {"aten::vstack", op::translate_vstack},
         {"aten::wait", op::skip_node},
         {"aten::where", op::translate_where},
         {"aten::zero", op::translate_zeros_like},

tests/layer_tests/pytorch_tests/test_hstack.py

@@ -0,0 +1,87 @@
+# Copyright (C) 2018-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import numpy as np
+import numpy as np
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+class aten_hstack(torch.nn.Module):
+    def forward(self, x):
+        return torch.hstack(self.prepare_input(x))
+    
+    def prepare_input(self, x):
+        return (x, x)
+    
+class aten_hstack_out(aten_hstack):
+    def forward(self, x, out):
+        return torch.hstack(self.prepare_input(x), out=out), out
+    
+class TestHstack(PytorchLayerTest):
+    def _prepare_input(self, out=False, num_repeats=2):
+        data = np.random.randn(2, 1, 3)
+        if not out:
+            return (data, )
+        concat = [data for _ in range(num_repeats)]
+        out = np.zeros_like(np.concatenate(concat, axis=1))
+        return (data, out)
+    
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("out", [False, True])
+    def test_hstack(self, out, ie_device, precision, ir_version):
+        model = aten_hstack() if not out else aten_hstack_out()
+        self._test(model, None, "aten::hstack", ie_device, 
+                   precision, ir_version, kwargs_to_prepare_input={"out": out, "num_repeats": 2})
+        
+    
+class TestHstackAlignTypes(PytorchLayerTest):
+    def _prepare_input(self, in_types):
+        in_vals = []
+        for i in range(len(in_types)):
+            dtype = in_types[i]
+            in_vals.append(np.random.randn(2, 1, 3).astype(dtype))
+        return in_vals
+    
+    def create_model(self, in_count):
+        class aten_align_types_hstack_two_args(torch.nn.Module):            
+            def forward(self, x, y):
+                ins = [x, y]
+                return torch.hstack(ins)
+        
+        class aten_align_types_hstack_three_args(torch.nn.Module):
+            def forward(self, x, y, z):
+                ins = [x, y, z]
+                return torch.hstack(ins)
+            
+        if in_count == 2:
+            return aten_align_types_hstack_two_args()
+        
+        if in_count == 3:
+            return aten_align_types_hstack_three_args()
+        
+    @pytest.mark.parametrize(("in_types"), [
+        (np.float32, np.int32),
+        (np.int32, np.float32),
+        (np.float16, np.float32),
+        (np.int16, np.float16),
+        (np.int32, np.int64),
+        # # Three inputs
+        (np.float32, np.int32, np.int32),
+        (np.float32, np.int32, np.float32),
+        (np.int32, np.float32, np.int32),
+        (np.float32, np.int32, np.int16),
+        (np.int32, np.float32, np.int16),
+        (np.int16, np.int32, np.int16),
+        (np.float16, np.float32, np.float16),
+        (np.float32, np.float16, np.float32),
+        (np.float16, np.int32, np.int16),
+        (np.int16, np.float16, np.int16)
+    ])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_align_types_hstack(self, ie_device, precision, ir_version, in_types):
+        self._test(self.create_model(len(in_types)), None, "aten::hstack",
+                   ie_device, precision, ir_version, kwargs_to_prepare_input={"in_types": in_types})

tests/layer_tests/pytorch_tests/test_vstack.py

@@ -0,0 +1,87 @@
+# Copyright (C) 2018-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+import numpy as np
+import numpy as np
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+class aten_vstack(torch.nn.Module):
+    def forward(self, x):
+        return torch.vstack(self.prepare_input(x))
+    
+    def prepare_input(self, x):
+        return (x, x)
+    
+class aten_vstack_out(aten_vstack):
+    def forward(self, x, out):
+        return torch.vstack(self.prepare_input(x), out=out), out
+    
+class TestVstack(PytorchLayerTest):
+    def _prepare_input(self, out=False, num_repeats=2):
+        data = np.random.randn(2, 1, 3)
+        if not out:
+            return (data, )
+        concat = [data for _ in range(num_repeats)]
+        out = np.zeros_like(np.concatenate(concat, axis=0))
+        return (data, out)
+    
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("out", [False, True])
+    def test_vstack(self, out, ie_device, precision, ir_version):
+        model = aten_vstack() if not out else aten_vstack_out()
+        self._test(model, None, "aten::vstack", ie_device, 
+                   precision, ir_version, kwargs_to_prepare_input={"out": out, "num_repeats": 2})
+        
+    
+class TestVstackAlignTypes(PytorchLayerTest):
+    def _prepare_input(self, in_types):
+        in_vals = []
+        for i in range(len(in_types)):
+            dtype = in_types[i]
+            in_vals.append(np.random.randn(2, 1, 3).astype(dtype))
+        return in_vals
+    
+    def create_model(self, in_count):
+        class aten_align_types_vstack_two_args(torch.nn.Module):
+            def forward(self, x, y):
+                ins = [x, y]
+                return torch.vstack(ins)
+        
+        class aten_align_types_vstack_three_args(torch.nn.Module):
+            def forward(self, x, y, z):
+                ins = [x, y, z]
+                return torch.vstack(ins)
+            
+        if in_count == 2:
+            return aten_align_types_vstack_two_args()
+        
+        if in_count == 3:
+            return aten_align_types_vstack_three_args()
+        
+    @pytest.mark.parametrize(("in_types"), [
+        (np.float32, np.int32),
+        (np.int32, np.float32),
+        (np.float16, np.float32),
+        (np.int16, np.float16),
+        (np.int32, np.int64),
+        # # Three inputs
+        (np.float32, np.int32, np.int32),
+        (np.float32, np.int32, np.float32),
+        (np.int32, np.float32, np.int32),
+        (np.float32, np.int32, np.int16),
+        (np.int32, np.float32, np.int16),
+        (np.int16, np.int32, np.int16),
+        (np.float16, np.float32, np.float16),
+        (np.float32, np.float16, np.float32),
+        (np.float16, np.int32, np.int16),
+        (np.int16, np.float16, np.int16)
+    ])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_align_types_vstack(self, ie_device, precision, ir_version, in_types):
+        self._test(self.create_model(len(in_types)), None, "aten::vstack",
+                   ie_device, precision, ir_version, kwargs_to_prepare_input={"in_types": in_types})