From 02e950cf548e5b5c8e92a67732284a6d42d7cc89 Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Wed, 25 Sep 2024 15:01:37 +0800
Subject: [PATCH 1/6] support inputs that have no batch

---
 src/frontends/pytorch/src/op/avg_poolnd.cpp   | 154 ++++++++++++-
 src/frontends/pytorch/src/op/max_poolnd.cpp   | 202 ++++++++++++++++++
 src/frontends/pytorch/src/op_table.cpp        |  34 +--
 .../pytorch_tests/test_avg_pool.py            | 157 ++++++++++++++
 .../pytorch_tests/test_max_pool.py            | 168 +++++++++++++++
 5 files changed, 699 insertions(+), 16 deletions(-)
 create mode 100644 tests/layer_tests/pytorch_tests/test_avg_pool.py
 create mode 100644 tests/layer_tests/pytorch_tests/test_max_pool.py

diff --git a/src/frontends/pytorch/src/op/avg_poolnd.cpp b/src/frontends/pytorch/src/op/avg_poolnd.cpp
index 03c32259b45091..ede81c90716c50 100644
--- a/src/frontends/pytorch/src/op/avg_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/avg_poolnd.cpp
@@ -8,9 +8,14 @@
 #include "openvino/op/concat.hpp"
 #include "openvino/op/constant.hpp"
 #include "openvino/op/pad.hpp"
-#include "openvino/op/subtract.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "openvino/op/slice.hpp"
+#include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/add.hpp"
 #include "utils.hpp"
 
+
 namespace ov {
 namespace frontend {
 namespace pytorch {
@@ -18,9 +23,8 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_avg_poolnd(const NodeContext& context) {
+OutputVector translate_avg_pool_base(const NodeContext& context, const Output<Node>& input) {
     num_inputs_check(context, 2, 7);
-    auto input = context.get_input(0);
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
     if (!context.input_is_none(2)) {
@@ -49,8 +53,152 @@ OutputVector translate_avg_poolnd(const NodeContext& context) {
                                 "Translation for aten::avg_pool2d do not support divisor_override input.");
     return {context.mark_node(
         std::make_shared<v14::AvgPool>(input, strides, pads, pads, kernel, !count_include_pad, rounding_type))};
+
+};
+
+OutputVector translate_avg_pool3d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-5})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_avg_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {5})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
+};
+
+OutputVector translate_avg_pool2d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-4})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+    
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_avg_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {4})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
+};
+
+OutputVector translate_avg_pool1d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+    
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_avg_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
 };
 
+
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp
index b6a01af1a7c2df..4b663c89b36753 100644
--- a/src/frontends/pytorch/src/op/max_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/max_poolnd.cpp
@@ -15,6 +15,10 @@
 #include "openvino/op/select.hpp"
 #include "openvino/op/shape_of.hpp"
 #include "openvino/op/subtract.hpp"
+#include "openvino/op/slice.hpp"
+#include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/add.hpp"
 #include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
@@ -72,11 +76,209 @@ OutputVector translate_max_poolnd(const NodeContext& context) {
     }
 };
 
+OutputVector translate_max_pool_base(const NodeContext& context, const Output<Node>& input) {
+    num_inputs_check(context, 3, 6);
+    auto kernel = context.const_input<Shape>(1);
+    Strides strides;
+    if (!context.input_is_none(2)) {
+        strides = context.const_input<Strides>(2);
+    }
+    if (context.input_is_none(2) || strides.size() == 0) {
+        // In case strides are not provided default is kernel
+        strides = kernel;
+    }
+    Shape pads;
+    if (context.input_is_none(3)) {
+        pads = Shape(kernel.size(), 0);
+    } else {
+        pads = context.const_input<Shape>(3);  // pytorch supports only symmetric paddings
+    }
+    Strides dilations;
+    if (!context.input_is_none(4)) {
+        dilations = context.const_input<Strides>(4);
+    }
+    RoundingType rounding_type;
+    if (context.input_is_none(5)) {
+        rounding_type = RoundingType::FLOOR;
+    } else {
+        rounding_type = context.const_input<bool>(5) ? RoundingType::CEIL_TORCH : RoundingType::FLOOR;
+    }
+
+    auto res = context.mark_node(std::make_shared<v14::MaxPool>(input,
+                                                                strides,
+                                                                dilations,
+                                                                pads,
+                                                                pads,
+                                                                kernel,
+                                                                rounding_type,
+                                                                PadType::EXPLICIT,
+                                                                element::i64,
+                                                                2));
+    if (context.get_output_size() == 2) {
+        auto out1 = res->output(0);
+        auto out2 = res->output(1);
+        return {std::move(out1), std::move(out2)};
+    } else {
+        return {res};
+    }
+};
+
+OutputVector translate_max_pool3d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-5})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_max_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {5})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
+};
+
+OutputVector translate_max_pool2d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-4})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_max_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {4})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
+};
+
+OutputVector translate_max_pool1d(const NodeContext& context) {
+    auto input = context.get_input(0);
+
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
+
+    // If there was no batch dimension, added dimension by slicing the input tensor
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
+
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
+    end_index, 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
+    auto pooled_output = translate_max_pool_base(context, input);
+
+    // If there was no batch dimension, remove the added dimension by slicing the output tensor
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
+
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {3})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
+
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+
+    for (auto& node : pooled_output) {
+        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
+    }
+    return pooled_output; 
+};
+
 OutputVector translate_max_poolnd_fx(const NodeContext& context) {
     auto output = translate_max_poolnd(context);
     return {context.mark_node(make_list_construct(output))};
 }
 
+OutputVector translate_max_pool2d_fx(const NodeContext& context) {
+    auto output = translate_max_pool2d(context);
+    return {context.mark_node(make_list_construct(output))};
+}
+
+OutputVector translate_max_pool3d_fx(const NodeContext& context) {
+    auto output = translate_max_pool3d(context);
+    return {context.mark_node(make_list_construct(output))};
+}
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp
index 1e4ecfc1e1367f..20c82ca02a2b9b 100644
--- a/src/frontends/pytorch/src/op_table.cpp
+++ b/src/frontends/pytorch/src/op_table.cpp
@@ -43,6 +43,9 @@ OP_CONVERTER(translate_argmin);
 OP_CONVERTER(translate_as_strided);
 OP_CONVERTER(translate_as_tensor);
 OP_CONVERTER(translate_avg_poolnd);
+OP_CONVERTER(translate_avg_pool1d);
+OP_CONVERTER(translate_avg_pool2d);
+OP_CONVERTER(translate_avg_pool3d);
 OP_CONVERTER(translate_bool);
 OP_CONVERTER(translate_batch_norm);
 OP_CONVERTER(translate_bitwise_and);
@@ -140,6 +143,9 @@ OP_CONVERTER(translate_masked_select);
 OP_CONVERTER(translate_max);
 OP_CONVERTER(translate_maximum);
 OP_CONVERTER(translate_max_poolnd);
+OP_CONVERTER(translate_max_pool1d);
+OP_CONVERTER(translate_max_pool2d);
+OP_CONVERTER(translate_max_pool3d);
 OP_CONVERTER(translate_mean);
 OP_CONVERTER(translate_meshgrid);
 OP_CONVERTER(translate_min);
@@ -282,6 +288,8 @@ OP_CONVERTER(translate_log_sigmoid_fx);
 OP_CONVERTER(translate_log_softmax_fx);
 OP_CONVERTER(translate_max_dim_fx);
 OP_CONVERTER(translate_max_poolnd_fx);
+OP_CONVERTER(translate_max_pool2d_fx);
+OP_CONVERTER(translate_max_pool3d_fx);
 OP_CONVERTER(translate_mean_fx);
 OP_CONVERTER(translate_min_dim_fx);
 OP_CONVERTER(translate_new_full_fx);
@@ -380,9 +388,9 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::atanh",
          op::optional_out<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Atanh>, 1>},
         {"aten::atanh_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Atanh>>},
-        {"aten::avg_pool1d", op::quantizable_op<op::translate_avg_poolnd>},
-        {"aten::avg_pool2d", op::quantizable_op<op::translate_avg_poolnd>},
-        {"aten::avg_pool3d", op::quantizable_op<op::translate_avg_poolnd>},
+        {"aten::avg_pool1d", op::quantizable_op<op::translate_avg_pool1d>},
+        {"aten::avg_pool2d", op::quantizable_op<op::translate_avg_pool2d>},
+        {"aten::avg_pool3d", op::quantizable_op<op::translate_avg_pool3d>},
         {"aten::baddbmm", op::translate_addmm},
         {"aten::batch_norm", op::translate_batch_norm},
         {"aten::bitwise_and", op::translate_bitwise_and},
@@ -534,12 +542,12 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::max", op::translate_max},
         {"aten::mv", op::translate_1to1_match_2_inputs<opset10::MatMul>},
         {"aten::maximum", op::translate_maximum},
-        {"aten::max_pool1d", op::quantizable_op<op::translate_max_poolnd>},
-        {"aten::max_pool1d_with_indices", op::quantizable_op<op::translate_max_poolnd>},
-        {"aten::max_pool2d", op::quantizable_op<op::translate_max_poolnd>},
-        {"aten::max_pool2d_with_indices", op::quantizable_op<op::translate_max_poolnd>},
-        {"aten::max_pool3d", op::quantizable_op<op::translate_max_poolnd>},
-        {"aten::max_pool3d_with_indices", op::quantizable_op<op::translate_max_poolnd>},
+        {"aten::max_pool1d", op::quantizable_op<op::translate_max_pool1d>},
+        {"aten::max_pool1d_with_indices", op::quantizable_op<op::translate_max_pool1d>},
+        {"aten::max_pool2d", op::quantizable_op<op::translate_max_pool2d>},
+        {"aten::max_pool2d_with_indices", op::quantizable_op<op::translate_max_pool2d>},
+        {"aten::max_pool3d", op::quantizable_op<op::translate_max_pool3d>},
+        {"aten::max_pool3d_with_indices", op::quantizable_op<op::translate_max_pool3d>},
         {"aten::mean", op::quantizable_op<op::translate_mean>},
         {"aten::meshgrid", op::translate_meshgrid},
         {"aten::min", op::translate_min},
@@ -768,8 +776,8 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_fx() {
         {"aten.asinh.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Asinh>},
         {"aten.atan.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Atan>},
         {"aten.atanh.default", op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Atanh>},
-        {"aten.avg_pool2d.default", op::translate_avg_poolnd},
-        {"aten.avg_pool3d.default", op::translate_avg_poolnd},
+        {"aten.avg_pool2d.default", op::translate_avg_pool2d},
+        {"aten.avg_pool3d.default", op::translate_avg_pool3d},
         {"aten.baddbmm.default", op::translate_addmm_fx},
         {"aten.bitwise_and.Scalar", op::translate_bitwise_and},
         {"aten.bitwise_and.Tensor", op::translate_bitwise_and},
@@ -866,8 +874,8 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_fx() {
         {"aten.masked_fill_.Tensor", op::inplace_op<op::translate_masked_fill>},
         {"aten.max.default", op::translate_max},
         {"aten.max.dim", op::translate_max_dim_fx},
-        {"aten.max_pool2d_with_indices.default", op::translate_max_poolnd_fx},
-        {"aten.max_pool3d_with_indices.default", op::translate_max_poolnd_fx},
+        {"aten.max_pool2d_with_indices.default", op::translate_max_pool2d_fx},
+        {"aten.max_pool3d_with_indices.default", op::translate_max_pool3d_fx},
         {"aten.maximum.default", op::translate_maximum},
         {"aten.mean.default", op::translate_mean_fx},
         {"aten.mean.dim", op::translate_mean_fx},
diff --git a/tests/layer_tests/pytorch_tests/test_avg_pool.py b/tests/layer_tests/pytorch_tests/test_avg_pool.py
new file mode 100644
index 00000000000000..af552299769869
--- /dev/null
+++ b/tests/layer_tests/pytorch_tests/test_avg_pool.py
@@ -0,0 +1,157 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import platform
+
+import pytest
+import torch
+
+from pytorch_layer_test_class import PytorchLayerTest
+import numpy as np
+
+
+d2_params = [{'kernel_size': [3, 3], 'stride': 1, 'padding': 0},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': 1},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [0, 1]},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [1, 0]},
+             {'kernel_size': [3, 3], 'stride': [2, 1], 'padding': 0},
+             {'kernel_size': [2, 1], 'stride': [2, 1], 'padding': 0},
+             {'kernel_size': [2, 1], 'stride': None, 'padding': 0},
+             {'kernel_size': [2, 1], 'stride': [], 'padding': 0},
+             ]
+
+d1_params = [{'kernel_size': 3, 'stride': 1, 'padding': 0},
+             {'kernel_size': (4,), 'stride': 1, 'padding': 1},
+             {'kernel_size': 4, 'stride': (5,), 'padding': 2},
+             {'kernel_size': 4, 'stride': None, 'padding': 0},
+             ]
+
+d3_params = [{'kernel_size': [3, 3, 3], 'stride': 1, 'padding': 0},
+             {'kernel_size': [3, 3, 3], 'stride': [1, 1, 1], 'padding': 1},
+             {'kernel_size': [3, 3, 3], 'stride': [3, 3, 3], 'padding': [0, 0, 0]},
+             {'kernel_size': [3, 2, 1], 'stride': [3, 1, 1], 'padding': [0, 0, 0]},
+             ]
+
+
+class TestAvgPool3D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
+        class aten_avg_pool3d(torch.nn.Module):
+
+            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                print(kernel_size)
+                self.stride = stride
+                print(stride)
+                self.padding = padding
+                self.ceil_mode = ceil_mode
+                self.count_include_pad = count_include_pad
+
+            def forward(self, input_tensor):
+                return torch.nn.functional.avg_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
+                                                      self.count_include_pad)
+        ref_net = None
+
+        return aten_avg_pool3d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool3d"
+
+
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
+    @pytest.mark.parametrize("params", d3_params)
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_avg_pool3d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
+                   ie_device, precision, ir_version, dynamic_shapes=True)
+    
+        
+        
+
+
+
+class TestAvgPool2D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
+        class aten_avg_pool2d(torch.nn.Module):
+
+            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                print(kernel_size)
+                self.stride = stride
+                print(stride)
+                self.padding = padding
+                self.ceil_mode = ceil_mode
+                self.count_include_pad = count_include_pad
+
+            def forward(self, input_tensor):
+                return torch.nn.functional.avg_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
+                                                      self.count_include_pad)
+        ref_net = None
+
+        return aten_avg_pool2d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool2d"
+
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15], [3, 15, 15]])    
+    @pytest.mark.parametrize("params", d2_params)
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_avg_pool2d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
+                   ie_device, precision, ir_version, trace_model=True)
+
+
+
+class TestAvgPool1D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
+        class aten_avg_pool1d(torch.nn.Module):
+
+            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                self.stride = stride
+                self.padding = padding
+                self.ceil_mode = ceil_mode
+                self.count_include_pad = count_include_pad
+
+            def forward(self, input_tensor):
+                return torch.nn.functional.avg_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
+                                                      self.count_include_pad)
+        ref_net = None
+
+        return aten_avg_pool1d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool1d"
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15], [3, 15]])
+    @pytest.mark.parametrize("params", d1_params)
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_avg_pool1d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
+                   ie_device, precision, ir_version, trace_model=True, dynamic_shapes=False)
+        
+
+
+
diff --git a/tests/layer_tests/pytorch_tests/test_max_pool.py b/tests/layer_tests/pytorch_tests/test_max_pool.py
new file mode 100644
index 00000000000000..2a6567549b3513
--- /dev/null
+++ b/tests/layer_tests/pytorch_tests/test_max_pool.py
@@ -0,0 +1,168 @@
+# Copyright (C) 2018-2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import platform
+
+import pytest
+import torch
+
+from pytorch_layer_test_class import PytorchLayerTest
+import numpy as np
+
+
+d2_params = [{'kernel_size': [3, 3], 'stride': 1, 'padding': 0},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': 1},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [0, 1]},
+             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [1, 0]},
+             {'kernel_size': [3, 3], 'stride': [2, 1], 'padding': 0},
+             {'kernel_size': [2, 1], 'stride': [2, 1], 'padding': 0},
+             {'kernel_size': [2, 1], 'stride': None, 'padding': 0},
+             ]
+
+d1_params = [{'kernel_size': 3, 'stride': 1, 'padding': 0},
+             {'kernel_size': (4,), 'stride': 1, 'padding': 1},
+             {'kernel_size': 4, 'stride': (5,), 'padding': 2},
+             {'kernel_size': 4, 'stride': None, 'padding': 0},
+             ]
+
+d3_params = [{'kernel_size': [3, 3, 3], 'stride': 1, 'padding': 0},
+             {'kernel_size': [3, 3, 3], 'stride': [1, 1, 1], 'padding': 1},
+             {'kernel_size': [3, 3, 3], 'stride': [3, 3, 3], 'padding': [0, 0, 0]},
+             {'kernel_size': [3, 2, 1], 'stride': [3, 1, 1], 'padding': [0, 0, 0]},
+             ]
+
+
+class TestMaxPool3D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
+        class aten_max_pool3d(torch.nn.Module):
+            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                self.stride = stride
+                self.padding = padding
+                self.dilation = dilation
+                self.ceil_mode = ceil_mode
+                self.dtype = dtype
+                self.return_indices = return_indices
+
+            def forward(self, input_tensor):
+                if self.return_indices :
+                    output, indices = torch.nn.functional.max_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
+                    return output, indices
+                else :
+                    output = torch.nn.functional.max_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
+                    return output, None
+        ref_net = None
+
+        if return_indices :
+            return aten_max_pool3d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool3d_with_indices"
+        return aten_max_pool3d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool3d"
+
+
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
+    @pytest.mark.parametrize("params", d3_params)
+    @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize('return_indices', [False, True])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_max_pool3d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
+                   ie_device, precision, ir_version, dynamic_shapes=True)
+        
+class TestMaxPool2D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
+        class aten_max_pool2d(torch.nn.Module):
+            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                self.stride = stride
+                self.padding = padding
+                self.dilation = dilation
+                self.ceil_mode = ceil_mode
+                self.dtype = dtype
+                self.return_indices = return_indices
+
+            def forward(self, input_tensor):
+                if self.return_indices :
+                    output, indices = torch.nn.functional.max_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
+                    return output, indices
+                else :
+                    output = torch.nn.functional.max_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
+                    return output, None
+        ref_net = None
+
+        if return_indices :
+            return aten_max_pool2d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool2d_with_indices"
+        return aten_max_pool2d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool2d"
+
+
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15], [3, 15, 15]])
+    @pytest.mark.parametrize("params", d2_params)
+    @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize('return_indices', [False, True])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_max_pool2d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
+                   ie_device, precision, ir_version, dynamic_shapes=True)
+        
+class TestMaxPool1D(PytorchLayerTest):
+    
+    def _prepare_input(self):
+        return (self.input_tensor, )
+    
+    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
+        class aten_max_pool1d(torch.nn.Module):
+            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
+                super().__init__()
+                self.kernel_size = kernel_size
+                self.stride = stride
+                self.padding = padding
+                self.dilation = dilation
+                self.ceil_mode = ceil_mode
+                self.dtype = dtype
+                self.return_indices = return_indices
+
+            def forward(self, input_tensor):
+                if self.return_indices :
+                    output, indices = torch.nn.functional.max_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
+                    return output, indices
+                else :
+                    output = torch.nn.functional.max_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
+                    return output, None
+        ref_net = None
+
+        if return_indices :
+            return aten_max_pool1d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool1d_with_indices"
+        return aten_max_pool1d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool1d"
+
+
+    @pytest.mark.parametrize('input_shape', [[1, 3, 15], [3, 15]])
+    @pytest.mark.parametrize("params", d1_params)
+    @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("ceil_mode", [True, False])
+    @pytest.mark.parametrize('return_indices', [False, True])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    def test_max_pool1d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
+                   ie_device, precision, ir_version, dynamic_shapes=True)

From cc6ed18c4a8bfcb462b9e08739fa8859e00f10fe Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Wed, 25 Sep 2024 20:39:10 +0800
Subject: [PATCH 2/6] fix code style

---
 src/frontends/pytorch/src/op/avg_poolnd.cpp | 174 +++------------
 src/frontends/pytorch/src/op/max_poolnd.cpp | 222 +++-----------------
 src/frontends/pytorch/src/op_table.cpp      |   3 -
 3 files changed, 69 insertions(+), 330 deletions(-)

diff --git a/src/frontends/pytorch/src/op/avg_poolnd.cpp b/src/frontends/pytorch/src/op/avg_poolnd.cpp
index ede81c90716c50..8130855a273bf8 100644
--- a/src/frontends/pytorch/src/op/avg_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/avg_poolnd.cpp
@@ -22,9 +22,25 @@ namespace pytorch {
 namespace op {
 
 using namespace ov::op;
-
-OutputVector translate_avg_pool_base(const NodeContext& context, const Output<Node>& input) {
+OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
     num_inputs_check(context, 2, 7);
+    auto input = context.get_input(0);
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+
+    auto const_0 = v0::Constant::create(element::i64, Shape{1}, {0});
+    auto const_1 = v0::Constant::create(element::i64, Shape{1}, {1});
+
+    // Unsqueeze axis to add batch dimension
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+
+    // Reshape pattern based on dimensions
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
+    auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
     if (!context.input_is_none(2)) {
@@ -51,154 +67,34 @@ OutputVector translate_avg_pool_base(const NodeContext& context, const Output<No
     }
     PYTORCH_OP_CONVERSION_CHECK(context.input_is_none(6),
                                 "Translation for aten::avg_pool2d do not support divisor_override input.");
-    return {context.mark_node(
-        std::make_shared<v14::AvgPool>(input, strides, pads, pads, kernel, !count_include_pad, rounding_type))};
+    auto res = context.mark_node(
+        std::make_shared<v14::AvgPool>(input, strides, pads, pads, kernel, !count_include_pad, rounding_type));
+
+    // Reshape back to original shape
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+    res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
+    return {res};
 
 };
 
-OutputVector translate_avg_pool3d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-5})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_avg_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {5})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+OutputVector translate_avg_pool1d(const NodeContext& context) {
+    return translate_avg_pool_base(context, 1);
 };
 
 OutputVector translate_avg_pool2d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-4})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-    
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_avg_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {4})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+    return translate_avg_pool_base(context, 2);
 };
 
-OutputVector translate_avg_pool1d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-    
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_avg_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+OutputVector translate_avg_pool3d(const NodeContext& context) {
+    return translate_avg_pool_base(context, 3);
 };
 
 
-
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp
index 4b663c89b36753..267a1b702732e5 100644
--- a/src/frontends/pytorch/src/op/max_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/max_poolnd.cpp
@@ -28,56 +28,25 @@ namespace pytorch {
 namespace op {
 
 using namespace ov::op;
-
-OutputVector translate_max_poolnd(const NodeContext& context) {
+OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
     num_inputs_check(context, 3, 6);
-    auto kernel = context.const_input<Shape>(1);
-    Strides strides;
-    if (!context.input_is_none(2)) {
-        strides = context.const_input<Strides>(2);
-    }
-    if (context.input_is_none(2) || strides.size() == 0) {
-        // In case strides are not provided default is kernel
-        strides = kernel;
-    }
-    Shape pads;
-    if (context.input_is_none(3)) {
-        pads = Shape(kernel.size(), 0);
-    } else {
-        pads = context.const_input<Shape>(3);  // pytorch supports only symmetric paddings
-    }
-    Strides dilations;
-    if (!context.input_is_none(4)) {
-        dilations = context.const_input<Strides>(4);
-    }
-    RoundingType rounding_type;
-    if (context.input_is_none(5)) {
-        rounding_type = RoundingType::FLOOR;
-    } else {
-        rounding_type = context.const_input<bool>(5) ? RoundingType::CEIL_TORCH : RoundingType::FLOOR;
-    }
+    auto input = context.get_input(0);
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
 
-    auto res = context.mark_node(std::make_shared<v14::MaxPool>(context.get_input(0),
-                                                                strides,
-                                                                dilations,
-                                                                pads,
-                                                                pads,
-                                                                kernel,
-                                                                rounding_type,
-                                                                PadType::EXPLICIT,
-                                                                element::i64,
-                                                                2));
-    if (context.get_output_size() == 2) {
-        auto out1 = res->output(0);
-        auto out2 = res->output(1);
-        return {std::move(out1), std::move(out2)};
-    } else {
-        return {res};
-    }
-};
+    auto const_0 = v0::Constant::create(element::i64, Shape{1}, {0});
+    auto const_1 = v0::Constant::create(element::i64, Shape{1}, {1});
+
+    // Unsqueeze axis to add batch dimension
+    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+
+    // Reshape pattern based on dimensions
+    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
+    auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
+    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
 
-OutputVector translate_max_pool_base(const NodeContext& context, const Output<Node>& input) {
-    num_inputs_check(context, 3, 6);
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
     if (!context.input_is_none(2)) {
@@ -114,170 +83,47 @@ OutputVector translate_max_pool_base(const NodeContext& context, const Output<No
                                                                 PadType::EXPLICIT,
                                                                 element::i64,
                                                                 2));
+
+    // Reshape back to original shape
+    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
+    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
+    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
+        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
     if (context.get_output_size() == 2) {
         auto out1 = res->output(0);
         auto out2 = res->output(1);
+        out1 = context.mark_node(std::make_shared<v1::Reshape>(out1, concat_shape, true));
+        out2 = context.mark_node(std::make_shared<v1::Reshape>(out2, concat_shape, true));
         return {std::move(out1), std::move(out2)};
     } else {
+        res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
         return {res};
     }
 };
 
-OutputVector translate_max_pool3d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-5})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_max_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {5})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+OutputVector translate_max_pool1d(const NodeContext& context) {
+    return translate_max_pool_base(context, 1);
 };
 
 OutputVector translate_max_pool2d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-4})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_max_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-2})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {4})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+    return translate_max_pool_base(context, 2);
 };
 
-OutputVector translate_max_pool1d(const NodeContext& context) {
-    auto input = context.get_input(0);
-
-    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-
-    auto unsqueeze_axis = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, unsqueeze_axis));
-
-    // If there was no batch dimension, added dimension by slicing the input tensor
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank,
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1}))));
-
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-3})), 
-    end_index, 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-    context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
-
-    auto pooled_output = translate_max_pool_base(context, input);
-
-    // If there was no batch dimension, remove the added dimension by slicing the output tensor
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(pooled_output[0]));
-
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {3})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {1})), 
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {0}))));
-
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-
-    for (auto& node : pooled_output) {
-        node = context.mark_node(std::make_shared<v1::Reshape>(node, concat_shape, true));
-    }
-    return pooled_output; 
+OutputVector translate_max_pool3d(const NodeContext& context) {
+    return translate_max_pool_base(context, 3);
 };
 
-OutputVector translate_max_poolnd_fx(const NodeContext& context) {
-    auto output = translate_max_poolnd(context);
-    return {context.mark_node(make_list_construct(output))};
-}
-
 OutputVector translate_max_pool2d_fx(const NodeContext& context) {
     auto output = translate_max_pool2d(context);
     return {context.mark_node(make_list_construct(output))};
-}
+};
 
 OutputVector translate_max_pool3d_fx(const NodeContext& context) {
     auto output = translate_max_pool3d(context);
     return {context.mark_node(make_list_construct(output))};
-}
+};
 
 }  // namespace op
 }  // namespace pytorch
diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp
index 20c82ca02a2b9b..dadcc4ff57bee1 100644
--- a/src/frontends/pytorch/src/op_table.cpp
+++ b/src/frontends/pytorch/src/op_table.cpp
@@ -42,7 +42,6 @@ OP_CONVERTER(translate_argmax);
 OP_CONVERTER(translate_argmin);
 OP_CONVERTER(translate_as_strided);
 OP_CONVERTER(translate_as_tensor);
-OP_CONVERTER(translate_avg_poolnd);
 OP_CONVERTER(translate_avg_pool1d);
 OP_CONVERTER(translate_avg_pool2d);
 OP_CONVERTER(translate_avg_pool3d);
@@ -142,7 +141,6 @@ OP_CONVERTER(translate_masked_scatter);
 OP_CONVERTER(translate_masked_select);
 OP_CONVERTER(translate_max);
 OP_CONVERTER(translate_maximum);
-OP_CONVERTER(translate_max_poolnd);
 OP_CONVERTER(translate_max_pool1d);
 OP_CONVERTER(translate_max_pool2d);
 OP_CONVERTER(translate_max_pool3d);
@@ -287,7 +285,6 @@ OP_CONVERTER(translate_leaky_relu_fx);
 OP_CONVERTER(translate_log_sigmoid_fx);
 OP_CONVERTER(translate_log_softmax_fx);
 OP_CONVERTER(translate_max_dim_fx);
-OP_CONVERTER(translate_max_poolnd_fx);
 OP_CONVERTER(translate_max_pool2d_fx);
 OP_CONVERTER(translate_max_pool3d_fx);
 OP_CONVERTER(translate_mean_fx);

From ce80a718f72a35ddf0c0e9cd1066b8e3e6d0af56 Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Thu, 26 Sep 2024 09:29:20 +0800
Subject: [PATCH 3/6] Different treatment of static and dynamic

---
 src/frontends/pytorch/src/op/avg_poolnd.cpp | 48 +++++++++-----
 src/frontends/pytorch/src/op/max_poolnd.cpp | 73 +++++++++++++--------
 2 files changed, 77 insertions(+), 44 deletions(-)

diff --git a/src/frontends/pytorch/src/op/avg_poolnd.cpp b/src/frontends/pytorch/src/op/avg_poolnd.cpp
index 8130855a273bf8..58c10c96424903 100644
--- a/src/frontends/pytorch/src/op/avg_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/avg_poolnd.cpp
@@ -11,6 +11,7 @@
 #include "openvino/op/shape_of.hpp"
 #include "openvino/op/slice.hpp"
 #include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/squeeze.hpp"
 #include "openvino/op/reshape.hpp"
 #include "openvino/op/add.hpp"
 #include "utils.hpp"
@@ -29,17 +30,22 @@ OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
 
     auto const_0 = v0::Constant::create(element::i64, Shape{1}, {0});
     auto const_1 = v0::Constant::create(element::i64, Shape{1}, {1});
+    bool is_static = input.get_partial_shape().rank().is_static();
+    bool no_batch_dim = is_static && input.get_partial_shape().rank().get_length() == dims + 1;
 
-    // Unsqueeze axis to add batch dimension
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
-
-    // Reshape pattern based on dimensions
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
-    auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+    if (is_static) {
+        if (no_batch_dim) {
+            input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+        }
+    } else {
+        input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+        auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+        auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+        auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
+        auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
+        auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+        input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+    }
 
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
@@ -70,14 +76,20 @@ OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
     auto res = context.mark_node(
         std::make_shared<v14::AvgPool>(input, strides, pads, pads, kernel, !count_include_pad, rounding_type));
 
-    // Reshape back to original shape
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-    res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
+    if (is_static) {
+        if (no_batch_dim) {
+            res = context.mark_node(std::make_shared<v0::Squeeze>(res, const_0));
+        }
+    } else {
+        auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
+        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
+            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
+        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
+            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+        auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+        res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
+    }
+    
     return {res};
 
 };
diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp
index 267a1b702732e5..40db8c3f7dbcc8 100644
--- a/src/frontends/pytorch/src/op/max_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/max_poolnd.cpp
@@ -17,6 +17,7 @@
 #include "openvino/op/subtract.hpp"
 #include "openvino/op/slice.hpp"
 #include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/squeeze.hpp"
 #include "openvino/op/reshape.hpp"
 #include "openvino/op/add.hpp"
 #include "openvino/op/util/framework_node.hpp"
@@ -35,17 +36,22 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
 
     auto const_0 = v0::Constant::create(element::i64, Shape{1}, {0});
     auto const_1 = v0::Constant::create(element::i64, Shape{1}, {1});
+    bool is_static = input.get_partial_shape().rank().is_static();
+    bool no_batch_dim = is_static && input.get_partial_shape().rank().get_length() == dims + 1;
 
-    // Unsqueeze axis to add batch dimension
-    input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
-
-    // Reshape pattern based on dimensions
-    auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
-    auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
-    auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
-    auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
-    auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
-    input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+    if (is_static) {
+        if (no_batch_dim) {
+            input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+        }
+    } else {
+        input = context.mark_node(std::make_shared<v0::Unsqueeze>(input, const_0));
+        auto unsqueeze_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input));
+        auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
+        auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
+        auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
+        auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+        input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
+    }
 
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
@@ -83,24 +89,39 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
                                                                 PadType::EXPLICIT,
                                                                 element::i64,
                                                                 2));
-
-    // Reshape back to original shape
-    auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
-    auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
-    auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
-        context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
-    auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
-    if (context.get_output_size() == 2) {
-        auto out1 = res->output(0);
-        auto out2 = res->output(1);
-        out1 = context.mark_node(std::make_shared<v1::Reshape>(out1, concat_shape, true));
-        out2 = context.mark_node(std::make_shared<v1::Reshape>(out2, concat_shape, true));
-        return {std::move(out1), std::move(out2)};
+    if (is_static) {
+        if (no_batch_dim) {
+            if (context.get_output_size() == 2) {
+                auto out1 = res->output(0);
+                auto out2 = res->output(1);
+                out1 = context.mark_node(std::make_shared<v0::Squeeze>(out1, const_0));
+                out2 = context.mark_node(std::make_shared<v0::Squeeze>(out2, const_0));
+                return {std::move(out1), std::move(out2)};
+            } else {
+                res = context.mark_node(std::make_shared<v0::Squeeze>(res, const_0));
+                return {res};
+            }
+        }
     } else {
-        res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
-        return {res};
+        auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
+        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
+            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
+        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
+            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+        auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+        if (context.get_output_size() == 2) {
+            auto out1 = res->output(0);
+            auto out2 = res->output(1);
+            out1 = context.mark_node(std::make_shared<v1::Reshape>(out1, concat_shape, true));
+            out2 = context.mark_node(std::make_shared<v1::Reshape>(out2, concat_shape, true));
+            return {std::move(out1), std::move(out2)};
+        } else {
+            res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
+            return {res};
+        }
     }
+
+    
 };
 
 OutputVector translate_max_pool1d(const NodeContext& context) {

From 9280c310bae68fdd64c6ebbd50a3ac21566288a7 Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Sun, 29 Sep 2024 22:38:49 +0800
Subject: [PATCH 4/6] fix code style

---
 src/frontends/pytorch/src/op/avg_poolnd.cpp | 12 +++++++----
 src/frontends/pytorch/src/op/max_poolnd.cpp | 23 ++++++++++++++++-----
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/frontends/pytorch/src/op/avg_poolnd.cpp b/src/frontends/pytorch/src/op/avg_poolnd.cpp
index 58c10c96424903..94de175aac5945 100644
--- a/src/frontends/pytorch/src/op/avg_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/avg_poolnd.cpp
@@ -82,10 +82,14 @@ OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
         }
     } else {
         auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
-        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
-            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
-        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
-            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+        
+        auto start_index_input = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
+        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
+
+        auto start_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
+        auto end_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims}));
+        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
+
         auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
         res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
     }
diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp
index 40db8c3f7dbcc8..ea65980a5327a0 100644
--- a/src/frontends/pytorch/src/op/max_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/max_poolnd.cpp
@@ -101,13 +101,26 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
                 res = context.mark_node(std::make_shared<v0::Squeeze>(res, const_0));
                 return {res};
             }
-        }
+        } else {
+            if (context.get_output_size() == 2) {
+                auto out1 = res->output(0);
+                auto out2 = res->output(1);
+                return {std::move(out1), std::move(out2)};
+            } else {
+                return {res};
+            }
+        } 
+
     } else {
         auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
-        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0,
-            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})), const_1, const_0));
-        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims})),
-            context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims})), const_1, const_0));
+
+        auto start_index_input = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
+        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
+
+        auto start_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
+        auto end_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims}));
+        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
+
         auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
         if (context.get_output_size() == 2) {
             auto out1 = res->output(0);

From da57218529f112b9b43a37935da9d4caaa1a5ba0 Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Thu, 10 Oct 2024 16:33:59 +0800
Subject: [PATCH 5/6] use test_pooling.py and return all case

---
 .../pytorch_tests/test_avg_pool.py            | 157 ----------------
 .../pytorch_tests/test_max_pool.py            | 168 ------------------
 .../layer_tests/pytorch_tests/test_pooling.py |  76 +++++---
 3 files changed, 50 insertions(+), 351 deletions(-)
 delete mode 100644 tests/layer_tests/pytorch_tests/test_avg_pool.py
 delete mode 100644 tests/layer_tests/pytorch_tests/test_max_pool.py

diff --git a/tests/layer_tests/pytorch_tests/test_avg_pool.py b/tests/layer_tests/pytorch_tests/test_avg_pool.py
deleted file mode 100644
index af552299769869..00000000000000
--- a/tests/layer_tests/pytorch_tests/test_avg_pool.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# Copyright (C) 2018-2024 Intel Corporation
-# SPDX-License-Identifier: Apache-2.0
-
-import platform
-
-import pytest
-import torch
-
-from pytorch_layer_test_class import PytorchLayerTest
-import numpy as np
-
-
-d2_params = [{'kernel_size': [3, 3], 'stride': 1, 'padding': 0},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': 1},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [0, 1]},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [1, 0]},
-             {'kernel_size': [3, 3], 'stride': [2, 1], 'padding': 0},
-             {'kernel_size': [2, 1], 'stride': [2, 1], 'padding': 0},
-             {'kernel_size': [2, 1], 'stride': None, 'padding': 0},
-             {'kernel_size': [2, 1], 'stride': [], 'padding': 0},
-             ]
-
-d1_params = [{'kernel_size': 3, 'stride': 1, 'padding': 0},
-             {'kernel_size': (4,), 'stride': 1, 'padding': 1},
-             {'kernel_size': 4, 'stride': (5,), 'padding': 2},
-             {'kernel_size': 4, 'stride': None, 'padding': 0},
-             ]
-
-d3_params = [{'kernel_size': [3, 3, 3], 'stride': 1, 'padding': 0},
-             {'kernel_size': [3, 3, 3], 'stride': [1, 1, 1], 'padding': 1},
-             {'kernel_size': [3, 3, 3], 'stride': [3, 3, 3], 'padding': [0, 0, 0]},
-             {'kernel_size': [3, 2, 1], 'stride': [3, 1, 1], 'padding': [0, 0, 0]},
-             ]
-
-
-class TestAvgPool3D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
-        class aten_avg_pool3d(torch.nn.Module):
-
-            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                print(kernel_size)
-                self.stride = stride
-                print(stride)
-                self.padding = padding
-                self.ceil_mode = ceil_mode
-                self.count_include_pad = count_include_pad
-
-            def forward(self, input_tensor):
-                return torch.nn.functional.avg_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
-                                                      self.count_include_pad)
-        ref_net = None
-
-        return aten_avg_pool3d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool3d"
-
-
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
-    @pytest.mark.parametrize("params", d3_params)
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize("count_include_pad", [True, False])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_avg_pool3d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, dynamic_shapes=True)
-    
-        
-        
-
-
-
-class TestAvgPool2D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
-        class aten_avg_pool2d(torch.nn.Module):
-
-            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                print(kernel_size)
-                self.stride = stride
-                print(stride)
-                self.padding = padding
-                self.ceil_mode = ceil_mode
-                self.count_include_pad = count_include_pad
-
-            def forward(self, input_tensor):
-                return torch.nn.functional.avg_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
-                                                      self.count_include_pad)
-        ref_net = None
-
-        return aten_avg_pool2d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool2d"
-
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15], [3, 15, 15]])    
-    @pytest.mark.parametrize("params", d2_params)
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize("count_include_pad", [True, False])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_avg_pool2d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, trace_model=True)
-
-
-
-class TestAvgPool1D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, ceil_mode=True, count_include_pad=True):
-        class aten_avg_pool1d(torch.nn.Module):
-
-            def __init__(self, kernel_size, stride, padding, ceil_mode, count_include_pad) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                self.stride = stride
-                self.padding = padding
-                self.ceil_mode = ceil_mode
-                self.count_include_pad = count_include_pad
-
-            def forward(self, input_tensor):
-                return torch.nn.functional.avg_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.ceil_mode,
-                                                      self.count_include_pad)
-        ref_net = None
-
-        return aten_avg_pool1d(kernel_size, stride, padding, ceil_mode=True, count_include_pad=True), ref_net, "aten::avg_pool1d"
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15], [3, 15]])
-    @pytest.mark.parametrize("params", d1_params)
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize("count_include_pad", [True, False])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_avg_pool1d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, trace_model=True, dynamic_shapes=False)
-        
-
-
-
diff --git a/tests/layer_tests/pytorch_tests/test_max_pool.py b/tests/layer_tests/pytorch_tests/test_max_pool.py
deleted file mode 100644
index 2a6567549b3513..00000000000000
--- a/tests/layer_tests/pytorch_tests/test_max_pool.py
+++ /dev/null
@@ -1,168 +0,0 @@
-# Copyright (C) 2018-2024 Intel Corporation
-# SPDX-License-Identifier: Apache-2.0
-
-import platform
-
-import pytest
-import torch
-
-from pytorch_layer_test_class import PytorchLayerTest
-import numpy as np
-
-
-d2_params = [{'kernel_size': [3, 3], 'stride': 1, 'padding': 0},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': 1},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [0, 1]},
-             {'kernel_size': [3, 3], 'stride': [1, 1], 'padding': [1, 0]},
-             {'kernel_size': [3, 3], 'stride': [2, 1], 'padding': 0},
-             {'kernel_size': [2, 1], 'stride': [2, 1], 'padding': 0},
-             {'kernel_size': [2, 1], 'stride': None, 'padding': 0},
-             ]
-
-d1_params = [{'kernel_size': 3, 'stride': 1, 'padding': 0},
-             {'kernel_size': (4,), 'stride': 1, 'padding': 1},
-             {'kernel_size': 4, 'stride': (5,), 'padding': 2},
-             {'kernel_size': 4, 'stride': None, 'padding': 0},
-             ]
-
-d3_params = [{'kernel_size': [3, 3, 3], 'stride': 1, 'padding': 0},
-             {'kernel_size': [3, 3, 3], 'stride': [1, 1, 1], 'padding': 1},
-             {'kernel_size': [3, 3, 3], 'stride': [3, 3, 3], 'padding': [0, 0, 0]},
-             {'kernel_size': [3, 2, 1], 'stride': [3, 1, 1], 'padding': [0, 0, 0]},
-             ]
-
-
-class TestMaxPool3D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
-        class aten_max_pool3d(torch.nn.Module):
-            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                self.stride = stride
-                self.padding = padding
-                self.dilation = dilation
-                self.ceil_mode = ceil_mode
-                self.dtype = dtype
-                self.return_indices = return_indices
-
-            def forward(self, input_tensor):
-                if self.return_indices :
-                    output, indices = torch.nn.functional.max_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
-                    return output, indices
-                else :
-                    output = torch.nn.functional.max_pool3d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
-                    return output, None
-        ref_net = None
-
-        if return_indices :
-            return aten_max_pool3d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool3d_with_indices"
-        return aten_max_pool3d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool3d"
-
-
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
-    @pytest.mark.parametrize("params", d3_params)
-    @pytest.mark.parametrize("dilation", [1, 2])
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize('return_indices', [False, True])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_max_pool3d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
-                   ie_device, precision, ir_version, dynamic_shapes=True)
-        
-class TestMaxPool2D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
-        class aten_max_pool2d(torch.nn.Module):
-            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                self.stride = stride
-                self.padding = padding
-                self.dilation = dilation
-                self.ceil_mode = ceil_mode
-                self.dtype = dtype
-                self.return_indices = return_indices
-
-            def forward(self, input_tensor):
-                if self.return_indices :
-                    output, indices = torch.nn.functional.max_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
-                    return output, indices
-                else :
-                    output = torch.nn.functional.max_pool2d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
-                    return output, None
-        ref_net = None
-
-        if return_indices :
-            return aten_max_pool2d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool2d_with_indices"
-        return aten_max_pool2d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool2d"
-
-
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15, 15], [3, 15, 15]])
-    @pytest.mark.parametrize("params", d2_params)
-    @pytest.mark.parametrize("dilation", [1, 2])
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize('return_indices', [False, True])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_max_pool2d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
-                   ie_device, precision, ir_version, dynamic_shapes=True)
-        
-class TestMaxPool1D(PytorchLayerTest):
-    
-    def _prepare_input(self):
-        return (self.input_tensor, )
-    
-    def create_model(self, kernel_size, stride, padding, dilation, ceil_mode=True, dtype=torch.float32, return_indices=False):
-        class aten_max_pool1d(torch.nn.Module):
-            def __init__(self, kernel_size, stride, padding, dilation, ceil_mode, return_indices) -> None:
-                super().__init__()
-                self.kernel_size = kernel_size
-                self.stride = stride
-                self.padding = padding
-                self.dilation = dilation
-                self.ceil_mode = ceil_mode
-                self.dtype = dtype
-                self.return_indices = return_indices
-
-            def forward(self, input_tensor):
-                if self.return_indices :
-                    output, indices = torch.nn.functional.max_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation,self.ceil_mode, True)
-                    return output, indices
-                else :
-                    output = torch.nn.functional.max_pool1d(input_tensor, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode, False)
-                    return output, None
-        ref_net = None
-
-        if return_indices :
-            return aten_max_pool1d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool1d_with_indices"
-        return aten_max_pool1d(kernel_size, stride, padding, dilation, ceil_mode, return_indices), ref_net, "aten::max_pool1d"
-
-
-    @pytest.mark.parametrize('input_shape', [[1, 3, 15], [3, 15]])
-    @pytest.mark.parametrize("params", d1_params)
-    @pytest.mark.parametrize("dilation", [1, 2])
-    @pytest.mark.parametrize("ceil_mode", [True, False])
-    @pytest.mark.parametrize('return_indices', [False, True])
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    @pytest.mark.precommit_torch_export
-    @pytest.mark.precommit_fx_backend
-    def test_max_pool1d(self, params, ceil_mode, dilation, return_indices, ie_device, precision, ir_version, input_shape):
-        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
-        self._test(*self.create_model(**params, ceil_mode=ceil_mode, dilation=dilation, return_indices=return_indices),
-                   ie_device, precision, ir_version, dynamic_shapes=True)
diff --git a/tests/layer_tests/pytorch_tests/test_pooling.py b/tests/layer_tests/pytorch_tests/test_pooling.py
index 32c8a973cb1c92..1924df2484f177 100644
--- a/tests/layer_tests/pytorch_tests/test_pooling.py
+++ b/tests/layer_tests/pytorch_tests/test_pooling.py
@@ -36,10 +36,8 @@
 
 
 class TestPooling(PytorchLayerTest):
-    def _prepare_input(self, ndim=4):
-        import numpy as np
-        shape = (1, 3, 15, 15, 15)
-        return (np.random.randn(*shape[:ndim]).astype(np.float32),)
+    def _prepare_input(self):
+        return (self.input_tensor,)
 
     def create_model(self, op_type, kernel_size, stride, padding, dilation=1, ceil_mode=True, count_include_pad=True, dtype=torch.float32):
         class aten_avg_pooling_base(torch.nn.Module):
@@ -129,121 +127,147 @@ def forward(self, x):
 
         return aten_pooling(), ref_net, f"aten::{op_type}"
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15], [3, 15]])
     @pytest.mark.parametrize("params", d1_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_torch_export
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_avg_pool1d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version):
+    def test_avg_pool1d(self, input_shape, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("avg_pool1d", **params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 3}, trace_model=True,
-                   dynamic_shapes=False)
+                   ie_device, precision, ir_version, trace_model=True,
+                   dynamic_shapes=is_dynamic_shapes)
 
-    @pytest.mark.parametrize(
-        "params", d2_params)
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15], [3, 15, 15]])
+    @pytest.mark.parametrize("params", d2_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_torch_export
     @pytest.mark.precommit_fx_backend
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_avg_pool2d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version):
+    def test_avg_pool2d(self, input_shape, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, is_dynamic_shapes):
         if ceil_mode and count_include_pad and np.array_equal(np.array(params["kernel_size"]), np.array([8, 8])):
             pytest.xfail("Ticket - 150292")
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("avg_pool2d", **params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, trace_model=True, freeze_model=False, dynamic_shapes=False)
+                   ie_device, precision, ir_version, trace_model=True, freeze_model=False, dynamic_shapes=is_dynamic_shapes)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
     @pytest.mark.parametrize("params", d3_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("count_include_pad", [True, False])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_torch_export
     @pytest.mark.precommit_fx_backend
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_avg_pool3d(self, params, ceil_mode, count_include_pad, ie_device, precision, ir_version):
+    def test_avg_pool3d(self, input_shape, params, ceil_mode, count_include_pad, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("avg_pool3d", **params, ceil_mode=ceil_mode, count_include_pad=count_include_pad),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 5}, trace_model=True,
-                   dynamic_shapes=False)
+                   ie_device, precision, ir_version, trace_model=True,
+                   dynamic_shapes=is_dynamic_shapes)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15], [3, 15]])
     @pytest.mark.parametrize("params", d1_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_torch_export
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool1d(self, params, ceil_mode, dilation, ie_device, precision, ir_version):
+    def test_max_pool1d(self, input_shape, params, ceil_mode, dilation, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool1d", **params, ceil_mode=ceil_mode, dilation=dilation),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 3}, dynamic_shapes=False)
+                   ie_device, precision, ir_version, dynamic_shapes=is_dynamic_shapes)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15], [3, 15, 15]])
     @pytest.mark.parametrize("params", d2_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
     @pytest.mark.parametrize("dtype", [torch.float32, torch.int32])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool2d(self, params, ceil_mode, dilation, dtype, ie_device, precision, ir_version):
+    def test_max_pool2d(self, input_shape, params, ceil_mode, dilation, dtype, ie_device, precision, ir_version, is_dynamic_shapes):
         to_trace = False
         if params["stride"] == []:
             to_trace = True
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool2d", **params, ceil_mode=ceil_mode, dilation=dilation, dtype=dtype),
-                   ie_device, precision, ir_version, dynamic_shapes=False, trace_model=to_trace)
+                   ie_device, precision, ir_version, dynamic_shapes=is_dynamic_shapes, trace_model=to_trace)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
     @pytest.mark.parametrize("params", d3_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool3d(self, params, ceil_mode, dilation, ie_device, precision, ir_version):
+    def test_max_pool3d(self, input_shape, params, ceil_mode, dilation, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool3d", **params, ceil_mode=ceil_mode, dilation=dilation),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 5}, dynamic_shapes=False)
+                   ie_device, precision, ir_version,  dynamic_shapes=is_dynamic_shapes)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15], [3, 15]])
     @pytest.mark.parametrize("params", d1_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool1d_indices(self, params, ceil_mode, dilation, ie_device, precision, ir_version):
+    def test_max_pool1d_indices(self, input_shape, params, ceil_mode, dilation, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool1d_with_indices", **params, ceil_mode=ceil_mode, dilation=dilation),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 3}, dynamic_shapes=False)
+                   ie_device, precision, ir_version, dynamic_shapes=is_dynamic_shapes)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15], [3, 15, 15]])
     @pytest.mark.parametrize("params", d2_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_fx_backend
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool2d_indices(self, params, ceil_mode, dilation,  ie_device, precision, ir_version):
+    def test_max_pool2d_indices(self, input_shape, params, ceil_mode, dilation,  ie_device, precision, ir_version, is_dynamic_shapes):
         to_trace = False
         if params["stride"] == []:
             to_trace = True
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool2d_with_indices", **params, ceil_mode=ceil_mode, dilation=dilation),
-                   ie_device, precision, ir_version, dynamic_shapes=False, trace_model=to_trace)
+                   ie_device, precision, ir_version, dynamic_shapes=is_dynamic_shapes, trace_model=to_trace)
 
+    @pytest.mark.parametrize("input_shape", [[1, 3, 15, 15, 15], [3, 15, 15, 15]])
     @pytest.mark.parametrize("params", d3_params)
     @pytest.mark.parametrize("ceil_mode", [True, False])
     @pytest.mark.parametrize("dilation", [1, 2])
+    @pytest.mark.parametrize("is_dynamic_shapes", [True, False])
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_fx_backend
     @pytest.mark.xfail(condition=platform.system() == 'Darwin' and platform.machine() == 'arm64',
                        reason='Ticket - 122715')
-    def test_max_pool3d_indices(self, params, ceil_mode, dilation, ie_device, precision, ir_version):
+    def test_max_pool3d_indices(self, input_shape, params, ceil_mode, dilation, ie_device, precision, ir_version, is_dynamic_shapes):
+        self.input_tensor = np.random.randn(*input_shape).astype(np.float32)
         self._test(*self.create_model("max_pool3d_with_indices", **params, ceil_mode=ceil_mode, dilation=dilation),
-                   ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 5}, dynamic_shapes=False)
+                   ie_device, precision, ir_version, dynamic_shapes=is_dynamic_shapes)

From f31d0c1796aee4d0dec138cbb08d67677b1341a8 Mon Sep 17 00:00:00 2001
From: tianyiSKY1 <tianyiZTY@icloud.com>
Date: Fri, 11 Oct 2024 17:12:04 +0800
Subject: [PATCH 6/6] fix code style

---
 src/frontends/pytorch/src/op/avg_poolnd.cpp | 25 +++++++++++----------
 src/frontends/pytorch/src/op/max_poolnd.cpp | 23 ++++++++++---------
 2 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/src/frontends/pytorch/src/op/avg_poolnd.cpp b/src/frontends/pytorch/src/op/avg_poolnd.cpp
index 94de175aac5945..d8223b04bfe690 100644
--- a/src/frontends/pytorch/src/op/avg_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/avg_poolnd.cpp
@@ -3,20 +3,19 @@
 //
 
 #include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/add.hpp"
 #include "openvino/op/avg_pool.hpp"
 #include "openvino/op/broadcast.hpp"
 #include "openvino/op/concat.hpp"
 #include "openvino/op/constant.hpp"
 #include "openvino/op/pad.hpp"
+#include "openvino/op/reshape.hpp"
 #include "openvino/op/shape_of.hpp"
 #include "openvino/op/slice.hpp"
-#include "openvino/op/unsqueeze.hpp"
 #include "openvino/op/squeeze.hpp"
-#include "openvino/op/reshape.hpp"
-#include "openvino/op/add.hpp"
+#include "openvino/op/unsqueeze.hpp"
 #include "utils.hpp"
 
-
 namespace ov {
 namespace frontend {
 namespace pytorch {
@@ -43,7 +42,8 @@ OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
         auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
         auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
         auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
-        auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+        auto reshape_pattern =
+            context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
         input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
     }
 
@@ -82,20 +82,22 @@ OutputVector translate_avg_pool_base(const NodeContext& context, int dims) {
         }
     } else {
         auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
-        
+
         auto start_index_input = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
-        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
+        auto slice_input_shape =
+            context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
 
         auto start_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
         auto end_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims}));
-        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
+        auto slice_pooled_output_shape = context.mark_node(
+            std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
 
-        auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+        auto concat_shape = context.mark_node(
+            std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
         res = context.mark_node(std::make_shared<v1::Reshape>(res, concat_shape, true));
     }
-    
-    return {res};
 
+    return {res};
 };
 
 OutputVector translate_avg_pool1d(const NodeContext& context) {
@@ -110,7 +112,6 @@ OutputVector translate_avg_pool3d(const NodeContext& context) {
     return translate_avg_pool_base(context, 3);
 };
 
-
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp
index ea65980a5327a0..b846de68d28b49 100644
--- a/src/frontends/pytorch/src/op/max_poolnd.cpp
+++ b/src/frontends/pytorch/src/op/max_poolnd.cpp
@@ -12,14 +12,13 @@
 #include "openvino/op/multiply.hpp"
 #include "openvino/op/pad.hpp"
 #include "openvino/op/range.hpp"
+#include "openvino/op/reshape.hpp"
 #include "openvino/op/select.hpp"
 #include "openvino/op/shape_of.hpp"
-#include "openvino/op/subtract.hpp"
 #include "openvino/op/slice.hpp"
-#include "openvino/op/unsqueeze.hpp"
 #include "openvino/op/squeeze.hpp"
-#include "openvino/op/reshape.hpp"
-#include "openvino/op/add.hpp"
+#include "openvino/op/subtract.hpp"
+#include "openvino/op/unsqueeze.hpp"
 #include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
@@ -49,7 +48,8 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
         auto rank = context.mark_node(std::make_shared<v0::ShapeOf>(unsqueeze_shape));
         auto end_index = context.mark_node(std::make_shared<v1::Add>(rank, const_1));
         auto start_index = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims - 2}));
-        auto reshape_pattern = context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
+        auto reshape_pattern =
+            context.mark_node(std::make_shared<v8::Slice>(unsqueeze_shape, start_index, end_index, const_1, const_0));
         input = context.mark_node(std::make_shared<v1::Reshape>(input, reshape_pattern, true));
     }
 
@@ -109,19 +109,22 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
             } else {
                 return {res};
             }
-        } 
+        }
 
     } else {
         auto pooled_output_shape = context.mark_node(std::make_shared<v3::ShapeOf>(res));
 
         auto start_index_input = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
-        auto slice_input_shape = context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
+        auto slice_input_shape =
+            context.mark_node(std::make_shared<v8::Slice>(input_shape, const_0, start_index_input, const_1, const_0));
 
         auto start_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-dims}));
         auto end_index_pooled = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {2 + dims}));
-        auto slice_pooled_output_shape = context.mark_node(std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
+        auto slice_pooled_output_shape = context.mark_node(
+            std::make_shared<v8::Slice>(pooled_output_shape, start_index_pooled, end_index_pooled, const_1, const_0));
 
-        auto concat_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
+        auto concat_shape = context.mark_node(
+            std::make_shared<v0::Concat>(OutputVector{slice_input_shape, slice_pooled_output_shape}, 0));
         if (context.get_output_size() == 2) {
             auto out1 = res->output(0);
             auto out2 = res->output(1);
@@ -133,8 +136,6 @@ OutputVector translate_max_pool_base(const NodeContext& context, int dims) {
             return {res};
         }
     }
-
-    
 };
 
 OutputVector translate_max_pool1d(const NodeContext& context) {