[PT FE] Support multiple FX operations (#23112)

### Details:
 - *Unpack nested tuples at model outputs*
 - *Support for `aten.cat.default` with default axis*
 - *Support `aten::reciprocal_` and `aten::abs_` for TorchScript case*
 - *Support:*
    *  `aten.adaptive_max_pool3d.default`
    *  `aten.adaptive_max_pool{N}d.default`
    *  `aten.amin.default`
    *  `aten.argmin.default`
    *  `aten.bitwise_{XXX}.default`
    *  `aten.clamp_max.default`
    *  `aten.clamp_{max/min}.Tensor`
    *  `aten.fill.Tensor`
    *  `aten.flip.default`
    *  `aten.fmod.{Scalar/Tensor}`
    *  `aten.ge.{Scalar/Tensor}`
    *  `aten.gt.Tensor`
    *  `aten.index_select.default`
    *  `aten.le.{Scalar/Tensor}`
    *  `aten.log{10/1p/2}.default`
    *  `aten.{max/maximum/mean/min/minimum}.default`
    *  `aten.min.dim`
    *  `aten.ne.Tensor`
    *  `aten.pow.{Scalar/Tensor_Tensor}`
    *  `aten.reciprocal.default`
    *  `aten.rsub.Tensor`
    *  `aten.scatter.value`
    *  `aten.select_scatter.default`
    *  `aten.sign.default`
    *  `aten.sqrt.default`
    *  `aten.sum.default`
    *  `aten.unfold.default`
    *  `aten.var.correction`
    *  *trigonometric ops*

 - *Enabled layer tests for added ops for torch.export case*

### Tickets:
 - *ticket-id*

Author: mvafin

.github/workflows/job_python_unit_tests.yml

@@ -231,7 +231,7 @@ jobs:
       - name: PyTorch torch.export Layer Tests
         if: ${{ fromJSON(inputs.affected-components).PyTorch_FE.test && runner.arch != 'ARM64' }} # Ticket: 126287
         run: |
-          python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/pytorch_tests -m precommit_torch_export --junitxml=${INSTALL_TEST_DIR}/TEST-pytorch.xml
+          python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/pytorch_tests -n logical -m precommit_torch_export --junitxml=${INSTALL_TEST_DIR}/TEST-pytorch.xml
         env:
           TEST_DEVICE: CPU
           TEST_PRECISION: FP32

src/bindings/python/src/openvino/frontend/pytorch/fx_decoder.py

@@ -7,7 +7,7 @@
 from openvino.frontend.pytorch.py_pytorch_frontend import _FrontEndPytorchDecoder as Decoder
 from openvino.frontend.pytorch.py_pytorch_frontend import _Type as DecoderType
 from openvino.runtime import op, PartialShape, Type as OVType, OVAny, Shape
-from openvino.frontend.pytorch.utils import maybe_convert_max_int, make_constant, fetch_attr, pt_to_ov_type_map
+from openvino.frontend.pytorch.utils import maybe_convert_max_int, make_constant, fetch_attr, pt_to_ov_type_map, torch_tensor_to_ov_const
 
 import torch
 
@@ -21,11 +21,11 @@ def __init__(self, pt_module, fx_gm, nodes=None, mark_node_callback=None, input_
         self.m_decoders = []
         self.pt_module = pt_module
         self.fx_gm = fx_gm
-        self.input_types = input_types
+        self.input_types = [OVAny(pt_to_ov_type_map[str(t)])
+                            for t in input_types]
         self.input_shapes = input_shapes
 
         self._input_signature = []
-        self._output_names = []
 
         if issubclass(type(pt_module), torch.fx.graph_module.GraphModule):
 
@@ -39,41 +39,36 @@ def __init__(self, pt_module, fx_gm, nodes=None, mark_node_callback=None, input_
                     self._input_signature.append(self._nodes[i].name)
                 elif self._nodes[i].op == 'output':
                     # Instead of putting output index, refer to its target
-                    args = self._nodes[i].args
-                    if isinstance(args[0], tuple):
-                        args = args[0]
-                    if isinstance(args[0], dict):
-                        for name, output in args[0].items():
-                            self._outputs.append(self._nodes.index(output))
-                            self._output_names.append(name)
-                    else:
-                        for output in args:
-                            self._outputs.append(self._nodes.index(output))
+                    uargs = self.unpack_containers(self._nodes[i].args)
+                    self._outputs = [(arg[0], self._nodes.index(arg[1])) for arg in uargs if arg[1] is not None]
 
         elif issubclass(type(pt_module), torch.fx.Node):
 
             self._nodes = nodes  # passed from outer context
 
             # FIXME: Quadratic complexity nodes*nodes considering the outer loop over all nodes
-            for i in range(len(self._nodes)):
-                if self._nodes[i] == pt_module:
-                    self._outputs = [i]
+            self._outputs = [("", self._nodes.index(pt_module))]
 
             # None in inputs mean the input is inlined or None (also considered inlined)
             self._inputs = [self._nodes.index(
                 arg) if arg in self._nodes else (arg,) for arg in pt_module.args]
 
             # FIXME: Find a better way to pass nested tuples to OV frontend. This is a temporary solution to flatten arguments.
             new_inputs = []
+            self.input_types = []
             for i in range(len(pt_module.args)):
-                if isinstance(pt_module.args[i], list) and any([isinstance(a, torch.fx.Node) for a in pt_module.args[i]]):
+                if isinstance(pt_module.args[i], (list, tuple)) and any([isinstance(a, torch.fx.Node) for a in pt_module.args[i]]):
                     for arg in pt_module.args[i]:
                         if arg in self._nodes:
                             new_inputs.append(self._nodes.index(arg))
                         else:
                             new_inputs.append((arg,))
+                        self.input_types.append(OVAny(DecoderType.List(
+                            TorchFXPythonDecoder.get_type_for_value(arg))))
                 else:
                     new_inputs.append(self._inputs[i])
+                    self.input_types.append(
+                        TorchFXPythonDecoder.get_type_for_value(self._inputs[i]))
             self._inputs = new_inputs
 
     def inputs(self):
@@ -83,6 +78,24 @@ def inputs(self):
     def is_input_inlined(self, index):
         return isinstance(self._inputs[index], tuple)
 
+    @staticmethod
+    def unpack_containers(arg):
+        if isinstance(arg, (tuple, list)):
+            res = []
+            for e in arg:
+                res.extend(TorchFXPythonDecoder.unpack_containers(e))
+            return res
+        elif isinstance(arg, dict):
+            res = []
+            for k, e in arg.items():
+                unpacked = TorchFXPythonDecoder.unpack_containers(e)
+                if len(unpacked) == 1:
+                    unpacked[0] = (k, unpacked[0][1])
+                res.extend(unpacked)
+            return res
+        else:
+            return [("", arg)]
+
     @staticmethod
     def arg_to_constant(arg):
         if isinstance(arg, list):
@@ -91,7 +104,7 @@ def arg_to_constant(arg):
                     arg[0]).__name__], Shape([len(arg)]), arg)
             else:
                 # TODO: which type should we use if list is empty? Need a signaling value here
-                return make_constant(int, Shape([0]), [])
+                return make_constant(OVType.i32, Shape([0]), [])
         elif isinstance(arg, bool):
             return make_constant(OVType.boolean, Shape([]), [arg])
         elif isinstance(arg, int):
@@ -103,10 +116,10 @@ def arg_to_constant(arg):
                 []), [arg])  # TODO: f32? why not f64?
         return None
 
-    
     def inlined_input(self, index):
         assert index < len(self._inputs), "Requested input doesn't exist"
-        assert isinstance(self._inputs[index], tuple), "Requested input which is not inlined"
+        assert isinstance(
+            self._inputs[index], tuple), "Requested input which is not inlined"
         assert self._inputs[index][0] is not None, "Requested None inlined input"
         constant = None
         arg = self._inputs[index][0]
@@ -144,13 +157,13 @@ def get_input_strides(self, index: int) -> list:
 
     def get_input_type(self, index):
         if index < len(self.input_types):
-            return OVAny(pt_to_ov_type_map[str(self.input_types[index])])
+            return self.input_types[index]
         input = self._raw_input(index)
         return self.get_type_for_value(input)
 
     def get_output_debug_name(self, index):
-        if self._output_names is not None and index < len(self._output_names):
-            return self._output_names[index]
+        if self._outputs is not None and index < len(self._outputs) and self._outputs[index][0]:
+            return self._outputs[index][0]
         name = getattr(self.pt_module, "name", "output")
         return name + ":" + str(index)
 
@@ -168,7 +181,8 @@ def get_shape_for_value(self, value):
                 return PartialShape(len(value.meta['tensor_meta'].shape) * [-1])
         return PartialShape.dynamic()
 
-    def get_type_for_value(self, value):
+    @staticmethod
+    def get_type_for_value(value):
         if issubclass(type(value), torch.fx.Node):
             if ('tensor_meta' in value.meta.keys()):
                 if value.meta['tensor_meta'] and isinstance(value.meta['tensor_meta'], torch.Tensor):
@@ -259,10 +273,10 @@ def get_schema(self):
         return self.pt_module.schema()
 
     def outputs(self):
-        return self._outputs
+        return [o[1] for o in self._outputs]
 
     def _raw_outputs(self):
-        return [self._nodes[x] for x in self._outputs]
+        return [self._nodes[x[1]] for x in self._outputs]
 
     def _raw_output(self, index):
         return self._raw_outputs()[index]
@@ -293,7 +307,7 @@ def as_constant(self):
         if self.pt_module.op == 'get_attr':
             # Extract Constant from FX module field
             ret = fetch_attr(self.fx_gm, self.pt_module.target)
-            ov_const = op.Constant(ret.numpy(force=True), shared_memory=True)
+            ov_const = torch_tensor_to_ov_const(ret, shared_memory=True)
             return ov_const.outputs()
 
         if not self.get_op_type() == 'prim::Constant':

src/frontends/pytorch/src/op/arange.cpp

@@ -85,72 +85,32 @@ OutputVector translate_arange(const NodeContext& context) {
 OutputVector translate_arange_fx(const NodeContext& context) {
     auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
     auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));
-    int dtype_port = -1;
     auto dtype = element::f32;
-    bool dtype_applied = false;
     auto num_inputs = context.get_input_size();
     ov::Output<Node> end;
-    ov::Output<Node> out_tensor;
     ov::Output<Node> start = zero;
     ov::Output<Node> step = one;
 
     if (num_inputs == 1) {
-        // aten::arange(Scalar end, tensor out)
+        // arange = torch.ops.aten.arange.default(_local_scalar_dense, dtype = torch.int8, device = device(type='cpu'),
+        // pin_memory = False);
         end = context.get_input(0);
-        out_tensor = end;  // context.input_is_none(1) ? end : context.get_input(1);
     } else if (num_inputs == 2) {
-        // aten::arange(Scalar end, tensor out)
         start = context.get_input(0);
         end = context.get_input(1);
-        out_tensor = end;  // context.input_is_none(1) ? end : context.get_input(1);
     } else if (num_inputs == 3) {
-        // aten::arange(Scalar start, Scalar end, Scalar step, Tensor out)
         start = context.get_input(0);
         end = context.get_input(1);
         step = context.get_input(2);
-        out_tensor = end;  // context.input_is_none(3) ? end : context.get_input(3);
-    } else if (num_inputs == 5) {
-        // aten::arange(Scalar end, ScalarType dtype, Layout, Device, bool pin_memory)
-        end = context.get_input(0);
-        out_tensor = end;
-        dtype_port = 1;
-    } else if (num_inputs == 6) {
-        // aten::arange(Scalar start, Scalar end, ScalarType dtype, Layout, Device, bool pin_memory)
-        start = context.get_input(0);
-        end = context.get_input(1);
-        out_tensor = end;
-        dtype_port = 2;
-        dtype_applied = true;
-    } else if (num_inputs == 7) {
-        // aten::arange(Scalar start, Scalar end, Scalar step, ScalarType dtype, Layout, Device, bool pin_memory)
-        start = context.get_input(0);
-        end = context.get_input(1);
-        step = context.get_input(2);
-        out_tensor = end;
-        dtype_port = 3;
-        dtype_applied = true;
     } else {
         PYTORCH_OP_CONVERSION_CHECK(false, "Not expected number of inputs for ", context.get_op_type());
     }
-    if (dtype_port >= 0 && !context.input_is_none(dtype_port)) {
-        if (std::dynamic_pointer_cast<v0::Constant>(
-                context.get_input_from_visible_context(dtype_port).get_node_shared_ptr())) {
-            dtype = convert_dtype(context.const_input<int64_t>(dtype_port));
-            dtype_applied = true;
-        } else if (const auto& fw_node =
-                       cast_fw_node(context.get_input(dtype_port).get_node_shared_ptr(), "prim::dtype")) {
-            out_tensor = fw_node->input_value(0);
-            dtype_applied = false;
-        } else {
-            PYTORCH_OP_CONVERSION_CHECK(false, "Couldn't get dtype input");
-        }
+    if (context.has_attribute("dtype")) {
+        dtype = context.get_attribute<element::Type>("dtype");
     }
-    auto r_end = context.mark_node(std::make_shared<v0::Convert>(end, dtype));
-    auto r_start = context.mark_node(std::make_shared<v0::Convert>(start, dtype));
-    auto r_step = context.mark_node(std::make_shared<v0::Convert>(step, dtype));
-    auto range = context.mark_node(std::make_shared<v4::Range>(r_start, r_end, r_step, dtype));
-    if (!dtype_applied) {
-        range = context.mark_node(std::make_shared<v1::ConvertLike>(range, out_tensor));
+    auto range = context.mark_node(std::make_shared<v4::Range>(start, end, step, dtype));
+    if (!context.has_attribute("dtype")) {
+        range = context.mark_node(std::make_shared<v1::ConvertLike>(range, context.get_input(0)));
     }
     return {range};
 };

src/frontends/pytorch/src/op/argmax_argmin.cpp

@@ -21,7 +21,7 @@ using namespace ov::op;
 
 namespace {
 OutputVector create_argmax_argmin_op(const NodeContext& context, TopKMode mode) {
-    num_inputs_check(context, 2, 3);
+    num_inputs_check(context, 1, 3);
     auto input = context.get_input(0);
     bool keep_dims = false;
     auto k = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));

src/frontends/pytorch/src/op/avg_poolnd.cpp

@@ -19,7 +19,7 @@ namespace op {
 using namespace ov::op;
 
 OutputVector translate_avg_poolnd(const NodeContext& context) {
-    num_inputs_check(context, 3, 7);
+    num_inputs_check(context, 2, 7);
     auto input = context.get_input(0);
     auto kernel = context.const_input<Shape>(1);
     Strides strides;

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

@@ -74,12 +74,18 @@ OutputVector translate_cat(const NodeContext& context) {
 
 OutputVector translate_cat_fx(const NodeContext& context) {
     // This translator is only needed to get axis as constant from external scope
-    num_inputs_check(context, 2, context.get_input_size());
+    num_inputs_check(context, 1, context.get_input_size());
     std::deque<Output<Node>> list_elems;
     for (size_t i = 0; i < context.get_input_size() - 1; i++) {
         list_elems.push_back(context.get_input(static_cast<int>(i)));
     }
-    auto axis = context.const_input<int64_t>(context.get_input_size() - 1);
+    int64_t axis = 0;
+    if (!context.get_input_type(context.get_input_size() - 1).is<type::List>()) {
+        // axis can be not present and that means that last input will have List type
+        axis = context.const_input<int64_t>(context.get_input_size() - 1);
+    } else {
+        list_elems.push_back(context.get_input(static_cast<int>(context.get_input_size() - 1)));
+    }
     return translate_cat_common(context, list_elems, axis, true);
 };
 

src/frontends/pytorch/src/op/cumsum.cpp

@@ -3,6 +3,7 @@
 //
 
 #include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/convert.hpp"
 #include "openvino/op/cum_sum.hpp"
 #include "utils.hpp"
 
@@ -28,6 +29,19 @@ OutputVector translate_cumsum(const NodeContext& context) {
     return {result};
 };
 
+OutputVector translate_cumsum_fx(const NodeContext& context) {
+    // cumsum = torch.ops.aten.cumsum.default(arg0_1, 0, dtype = torch.float64)
+    num_inputs_check(context, 2, 2);
+    auto x = context.get_input(0);
+    auto dim = context.get_input(1);
+    if (context.has_attribute("dtype")) {
+        auto dtype = context.get_attribute<element::Type>("dtype");
+        x = context.mark_node(std::make_shared<v0::Convert>(x, dtype));
+    }
+    auto result = context.mark_node(std::make_shared<v0::CumSum>(x, dim));
+    return {result};
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/op/mean.cpp

@@ -3,6 +3,7 @@
 //
 
 #include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/convert.hpp"
 #include "openvino/op/reduce_mean.hpp"
 #include "utils.hpp"
 
@@ -11,6 +12,8 @@ namespace frontend {
 namespace pytorch {
 namespace op {
 
+using namespace ov::op;
+
 OutputVector translate_mean(const NodeContext& context) {
     num_inputs_check(context, 2, 5);
     auto x = context.get_input(0);
@@ -35,20 +38,26 @@ OutputVector translate_mean(const NodeContext& context) {
             x = apply_dtype(context, 3, x);
         }
     }
-    auto mean = context.mark_node(std::make_shared<ov::op::v1::ReduceMean>(x, axes, keep_dims));
+    auto mean = context.mark_node(std::make_shared<v1::ReduceMean>(x, axes, keep_dims));
     if (num_inputs == 5 && !context.input_is_none(4)) {
         context.mutate_input(4, mean);
     }
     return {mean};
 };
 
 OutputVector translate_mean_fx(const NodeContext& context) {
-    num_inputs_check(context, 2, 5);
+    num_inputs_check(context, 1, 5);
     auto x = context.get_input(0);
     auto num_inputs = context.get_input_size();
     bool keep_dims = false;
+    if (context.has_attribute("dtype")) {
+        auto dtype = context.get_attribute<element::Type>("dtype");
+        x = context.mark_node(std::make_shared<v0::Convert>(x, dtype));
+    }
     Output<Node> axes;
-    if (num_inputs == 2) {
+    if (num_inputs == 1) {
+        axes = get_node_axes_range(context, x);
+    } else if (num_inputs == 2) {
         axes = context.get_input(1);
     } else {
         axes = context.get_input(1);
@@ -59,7 +68,7 @@ OutputVector translate_mean_fx(const NodeContext& context) {
             x = apply_dtype(context, 3, x);
         }
     }
-    auto mean = context.mark_node(std::make_shared<ov::op::v1::ReduceMean>(x, axes, keep_dims));
+    auto mean = context.mark_node(std::make_shared<v1::ReduceMean>(x, axes, keep_dims));
     if (num_inputs == 5 && !context.input_is_none(4)) {
         context.mutate_input(4, mean);
     }