Torchfx backend (#18244)

* Added Torchscript Backend

* First commit for backend with Torch FX Decoder

* Merging changes from Torch FX branch

* Torch FX initial fixes (Temporary)

* Fixed type/shape issues in Torch FX decoder

* Added translation for built-in getitem

* MaxPool update & Output shape fix (Torch FX)

* Torch FX graph outputs fix

* Torch FX support for sigmoid and slu_

* Torch FX graph module caching

* Torch Fx partitioner cache removed

* Torch FX initial getitem replacer added

* Index check for torch fx getitem replacer

* Debug print removed from partitioner

* Added environment variables for pytorch tracing mode and openvino device

* FX translation fix for getitem & getitem replacer removed

* Added checks for PyTorch tracing mode environment variable

* Adding compile mode for fallback

* Added more ops for resnet18

* Added a check for environment variable

* Generalized addmm to work with torchscript and torchfx

* Added the missing batch_norm.default translation

* fx_backend: include get_attr ops to the partitions

* AddeTODO note t to improvget_attr algorithm

* created function for adding get_attr nodes

* fx_backend: added aten.mul.Tensor, re-enabled aten.empty.memory_format

* fx_backend: Additional op support/improvement for Inception V3

* Added comment for fix 64-bit to 32-bit max int conversion

* fx_backend: Update for avg_poolnd to support 3 inputs

* Fixed erorr in decoder.py

* TorchFX caching fix

* Torch backend, op support for Stable Diff. & BERT

* Arranged ops in order and added torch tensor mapping

* Added support for more ops for super glue

* TorchFX: Initial permanent fallback

* TorchFX: New ops for improved TorchVision support

* TorchFX backend optimizations for partitioning and tmp fallback

* working operator updates for superglue

* Updates to operators for superglue

* Removed max.dim and stack

* Cleanup

* Cleanup

* Fixed a couple of syntax issues

* Fixed a couple of syntax issues

* Added missing method to TorchFX Decoder

* Added missing method to TorchFX Decoder

* Removed redundant code for transpose

* TorchFX: Initial StableDiffusion support

* PyTorch decoder ovtype to ctype fix for int64

* Added ops for distilbert

* Fixed few unnecessary include statements

* Seperated TorchFX and TorchScript decoders

* Modified import statements to reflect two decoders

* f64 fix for TorchFX

* Import fix for PyTorch backend modules

* TorchFX serialize graph for debugging (Temporary)

* Serialize and load back feature enabled for TorchFX

* Temporary optimization to remove Broadcast

* Temporary SoftmaxRehapeElimination pass is added

* TorchFX custom model cache directory

* PyTorch bitwise translation, conversion checks enabled

* Naming fix in make_list_construct

* TorchFX: Added comments to Softmax and Slice translations

* translate_chunk temporarily removed for TS backend

* Fixed linter issues

* Addressed clang formatting issues

* Fixed few more clang and linter issues

* Fixed tests to use ts_decoder

* Fixed naming convention issues

* Added missing import

* Added inlined_inputs to TorchScriptDecoder

* Added tests for torch fx backend

* Removed magic numbers in PyTorch decoder utils

* TorchFX decoder data type fix

* Added cast from size_t to int

* TorchFX output handling code cleanup

* TorchFX: Use detached input tensor

* Added missing cast from size_t to int

* Added static cast in group_norm

* Fixed casting issue in split

---------

Co-authored-by: ynimmaga <yamini.nimmagadda@intel.com>
Co-authored-by: Cavus Mustafa <mustafa.cavus@intel.com>

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

@@ -0,0 +1,403 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+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, ov_to_c_type_map
+
+import typing
+from packaging.version import parse
+import torch
+import numpy as np
+import inspect
+import ctypes
+
+class TorchFXPythonDecoder (Decoder):
+
+    def __init__(self, pt_module, fx_gm, nodes=None, mark_node_callback=None, input_shapes=[], input_types=[]):
+        Decoder.__init__(self)
+        self.mark_node_callback = mark_node_callback
+        # We store every decoder created by this decoder so that all them are not deleted until the first decoder is deleted
+        self.m_decoders = []
+        self.pt_module = pt_module
+        self.fx_gm = fx_gm
+        self.input_types = input_types
+        self.input_shapes = input_shapes
+
+        if issubclass(type(pt_module), torch.fx.graph_module.GraphModule):
+
+            self._nodes = list(pt_module.graph.nodes)
+            self._inputs = []
+            self._outputs = []
+            for i in range(len(self._nodes)):
+                if self._nodes[i].op == 'placeholder':
+                    self._inputs.append(i)
+                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]
+                    for output in args:
+                        self._outputs.append(self._nodes.index(output))
+
+        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]
+
+            # code constant or None input as a tuple to diffirentiate it from regular input
+            # negative index is used to mark inputs initialized by inline constants, there are no such inputs in the graph
+            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 temprary solution to flatten arguments.
+            new_inputs = []
+            for i in range(len(pt_module.args)):
+                expand_list = False
+                if isinstance(pt_module.args[i], list):
+                    for arg in pt_module.args[i]:
+                        if arg in self._nodes:
+                            expand_list = True
+                            break;
+                if expand_list:
+                    for arg in pt_module.args[i]:
+                        new_inputs.append(self._nodes.index(arg))
+                else:
+                    new_inputs.append(self._inputs[i])
+            self._inputs = new_inputs
+
+
+
+    def inputs(self):
+        return [x if x is not None else 100000 for x in self._inputs]
+
+    def input(self, index):  # TODO: remove
+        return self.inputs()[index]  # TODO: find specialized method
+
+    def get_input_debug_name(self, index):
+        return "input"+str(index)
+
+    def get_input_signature_name(self, index: int) -> str:
+        return self.get_input_debug_name(index)
+
+    def get_input_shape(self, index):
+        if index < len(self.input_shapes):
+            return PartialShape(self.input_shapes[index])
+        input = self._raw_input(index)
+        return self.get_shape_for_value(input)
+
+    def get_input_type(self, index):
+        if index < len(self.input_types):
+            return OVAny(pt_to_ov_type_map[str(self.input_types[index])])
+        input = self._raw_input(index)
+        return self.get_type_for_value(input)
+
+    def get_output_debug_name(self, index):
+        return "output"+str(index)
+
+    def get_output_shape(self, index):
+        output = self._raw_output(index)
+        return self.get_shape_for_value(output)
+
+    def get_output_type(self, index):
+        output = self._raw_output(index)
+        return self.get_type_for_value(output)
+
+    def _get_known_type_for_value(self, type):
+        '''
+            Returns known/unknown types wrapped as OVAny
+        '''
+        # Check for simple scalar types first
+        # TODO: Don't use str, use native types
+        if type is None:
+            return OVAny(OVType.dynamic)
+        if str(type) in pt_to_ov_type_map:
+            return OVAny(pt_to_ov_type_map[str(type)])
+        elif type.__class__ is torch.TensorType:
+            # Tensor type, parse element type
+            # TODO: replace string by native type
+            # return OVAny(PartialShape([1,2,3]))
+            return OVAny(DecoderType.Tensor(self._get_known_type_for_value(type.dtype())))
+        elif type.__class__ is torch.ListType:
+            element_type = type.getElementType()
+            return OVAny(DecoderType.List(self._get_known_type_for_value(element_type)))
+        else:
+            # Not yet recognized
+            return OVAny(OVType.dynamic)
+            #pt_type_class = value.type().__class__
+            #    if pt_type_class is torch.ListType:
+
+    def get_shape_for_value(self, value):
+        if value and ('tensor_meta' in value.meta.keys()):
+            return PartialShape(value.meta['tensor_meta'].shape)
+        return PartialShape([1])
+
+    def get_type_for_value(self, value):
+        if issubclass(type(value), torch.fx.Node):
+            if ('tensor_meta' in value.meta.keys()):
+                pt_type = value.meta['tensor_meta'].dtype
+                if str(pt_type) in pt_to_ov_type_map:
+                    ov_type = pt_to_ov_type_map[str(pt_type)]
+                    return OVAny(ov_type)
+            else:
+                return OVAny(OVType.f32)
+        elif isinstance(value, int):
+            return OVAny(OVType.i32)
+        elif isinstance(value, float):
+            return OVAny(OVType.f32)
+        elif isinstance(value, bool):
+            return OVAny(OVType.boolean)
+        else:
+            return OVAny(OVType.f32)
+
+    def get_input_transpose_order(self, index):
+        return []
+        # TODO TBD
+
+        input = self._raw_input(index)
+        if input.type() is not None and input.type().kind() == 'TensorType':
+            strides = input.type().strides()
+            if strides is not None:
+                return [s[0] for s in sorted(enumerate(strides), key=lambda x:x[1], reverse=True)]
+        return []
+
+    def get_output_transpose_order(self, index):
+        return []
+
+        # old code
+        output = self._raw_output(index)
+        if output.type() is not None and output.type().kind() == 'TensorType':
+            strides = output.type().strides()
+            if strides is not None:
+                return [s[0] for s in sorted(enumerate(strides), key=lambda x:x[1], reverse=True)]
+        return []
+
+    def get_subgraph_size(self):
+        if issubclass(type(self.pt_module), torch.fx.Node):
+            return 0
+        return len(self.get_subgraphs()) if hasattr(self.pt_module, 'blocks') else 1
+
+    def visit_subgraph(self, node_visitor):
+        # make sure topological order is satisfied
+        for node in self._nodes:
+            if node.op == 'placeholder' or node.op == 'output':
+                continue # skipping non-operational nodes
+            decoder = TorchFXPythonDecoder(node, self.fx_gm, self._nodes, mark_node_callback=self.mark_node_callback)
+            self.m_decoders.append(decoder)
+            node_visitor(decoder)
+
+    def get_subgraphs(self):
+        if issubclass(type(self.pt_module), torch.fx.Node):
+            return []
+        return list(self.pt_module.blocks())
+
+    def get_subgraph_decoder(self, index):
+        decoder = TorchFXPythonDecoder(self.get_subgraphs()[index], self.fx_gm, mark_node_callback=self.mark_node_callback)
+        self.m_decoders.append(decoder)
+        return decoder
+
+    def get_op_type(self):
+        if self.pt_module.op == 'call_function':
+            return str(self.pt_module.target)
+        elif self.pt_module.op == 'get_attr':
+            return 'get_attr'  # FIXME should be aligned with get_attr from TS implementation
+        else:
+            return 'UNKNOWN_TYPE_' + str(self.pt_module.op)
+
+    def get_schema(self):
+        return ''
+        return self.pt_module.schema()
+
+    def outputs(self):
+        return self._outputs
+
+    def _raw_outputs(self):
+        return [self._nodes[x] for x in self._outputs]
+
+    def _raw_output(self, index):
+        return self._raw_outputs()[index]
+
+    def _raw_inputs(self):
+        return [self._nodes[x] if x is not None and x < len(self._nodes) else x for x in self._inputs]
+
+    def _raw_input(self, index):
+        return self._raw_inputs()[index]
+
+    def num_of_outputs(self):
+        return len(self.outputs())
+
+    def output(self, index):
+        return self.outputs()[index]
+
+    def mark_node(self, node):
+        if self.mark_node_callback is not None:
+            self.mark_node_callback(self, node)
+        return node
+
+    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)
+            ovshape = PartialShape(ret.size())
+            ovtype = pt_to_ov_type_map[str(ret.type())]
+            c_type = ctypes.POINTER(ov_to_c_type_map[ovtype])
+            data_c_ptr = ctypes.cast(ret.data_ptr(), c_type)
+            ov_const = op.Constant(ovtype, ovshape.get_shape(), data_c_ptr[:ret.nelement()])
+            return ov_const.outputs()
+
+
+        if not self.get_op_type() == 'prim::Constant':
+            return None
+        pt_value = self._raw_output(0)
+
+        pt_type_class = pt_value.type().__class__
+        if pt_type_class is torch.TensorType:
+            return self.as_constant_tensor(pt_value)
+        if pt_type_class is torch.ListType:
+            return self.as_constant_list(pt_value)
+        if str(pt_value.type()) in ['torch.int32', 'int']:
+            return make_constant(OVType.i32, Shape([]), [pt_value.toIValue()]).outputs()
+        if str(pt_value.type()) in ['torch.float', 'torch.FloatType', 'float']:
+            return make_constant(OVType.f32, Shape([]), [pt_value.toIValue()]).outputs()
+        if str(pt_value.type()) in ['torch.bool', 'bool']:
+            return make_constant(OVType.boolean, Shape([]), [pt_value.toIValue()]).outputs()
+
+        return None
+
+    def as_string(self):
+        if not self.get_op_type() == 'prim::Constant':
+            return None
+        pt_value = self._raw_output(0)
+
+        if str(pt_value.type()) in ['torch.StringType', 'str']:
+            return pt_value.toIValue()
+        return None
+
+    def as_constant_tensor(self, pt_value):
+        ivalue = pt_value.toIValue()
+        if pt_value.isCompleteTensor():
+            try:
+                ivalue = ivalue.to(memory_format=torch.contiguous_format).detach().cpu()
+            except:
+                print("[ WARNING ] Tensor couldn't detach")
+            if str(pt_value.type().dtype()) in pt_to_py_type_map:
+                # Constant interpretation doesn't respect new-full type of PT
+                # It recognizes only tensors, and give lists as 1D tensors, and scalars as Tensor scalars
+                # So only tensor-type constants are supported
+                ovshape = PartialShape(pt_value.type().sizes())
+                ovtype = pt_to_ov_type_map[str(pt_value.type().dtype())]
+
+                # TODO: try-except here is a temporary WA for issues with data_ptr that we currently cannot predict; provide better solution
+                try:
+                    # this is only possible with adding a new ctor for Constant Python binding
+                    # TODO Check strides and pass them somehow
+                    values = ivalue.data_ptr()
+                    ov_const = make_constant(ovtype, ovshape.get_shape(), values)
+                except:
+                    # old variant that makes a slow data copying
+                    print(f"[ WARNING ] Constant wasn't able to convert from data_ptr.")
+                    values = ivalue.flatten().tolist()
+                    ov_const = make_constant(ovtype, ovshape.get_shape(), values)
+                return ov_const.outputs()
+        else:
+            # Incomplete tensor can be scalar
+            if isinstance(ivalue, float):
+                return make_constant(OVType.f32, Shape([]), [ivalue]).outputs()
+            if isinstance(ivalue, int):
+                return make_constant(OVType.i32, Shape([]), [ivalue]).outputs()
+            if isinstance(ivalue, bool):
+                return make_constant(OVType.boolean, Shape([]), [ivalue]).outputs()
+
+            # TODO: verify that it correctly reads incomplete consts
+            if str(ivalue.type()) in pt_to_ov_type_map:
+                try:
+                    ovshape = PartialShape(ivalue.size())
+                    ovtype = pt_to_ov_type_map[str(ivalue.type())]
+                    ov_const = make_constant(ovtype, ovshape.get_shape(), ivalue.data_ptr())
+                except:
+                    # old variant that makes a slow data copying
+                    print(f"[ WARNING ] Constant wasn't able to convert from data_ptr.")
+                    nvalues = ivalue.numpy()
+                    ovtype = np_to_ov_type_map[str(nvalues.dtype)]
+                    ovshape = PartialShape(nvalues.shape)
+                    ov_const = make_constant(ovtype, ovshape.get_shape(), nvalues.flatten().tolist())
+                return ov_const.outputs()
+        return None
+
+    def as_constant_list(self, pt_value):
+        # For now it is treat a list as a 1D tensor; it is required by converters to avoid need to massively rewrite them in that part where constant attributes are queried
+        pt_element_type = str(pt_value.type().getElementType())
+        ivalue = pt_value.toIValue()
+        is_known_type = pt_element_type in pt_to_ov_type_map
+
+        # WA to broken ov.Type
+        # Detect integer list and process it with a dedicated method
+        # TODO: Fix ov.Type and remove this WA
+        # if pt_to_py_type_map[pt_element_type] == 'int':
+        #    self.as_constant_list_of_ints(ovshape = PartialShape([len(ivalue)]), ivalue)
+        # End of WA to broken ov.Type
+
+        if is_known_type:
+            ovtype = pt_to_ov_type_map[pt_element_type]
+            ovshape = PartialShape([len(ivalue)])
+            ov_const = make_constant(ovtype, ovshape.get_shape(), ivalue)
+            return ov_const.outputs()
+
+    def input_is_none(self, index):
+        if index >= len(self.inputs()) or self._raw_input(index) is None:
+            return True
+        else:
+            r_input = self._raw_input(index)
+            return str(type(r_input)) in ['torch.NoneType', 'NoneType']
+
+    def debug(self):
+        self.pt_module.print()
+
+    def inlined_inputs(self, index):
+        result = []
+        for i in range(len(self._inputs)):
+            if isinstance(self._inputs[i], tuple):
+                constant = None
+                arg = self._inputs[i][0]
+                if isinstance(arg, list):
+                    if len(arg) > 0:
+                        constant = make_constant(pt_to_ov_type_map[type(arg[0]).__name__], Shape([len(arg)]), arg)
+                    else:
+                        # TODO: which type should we use if list is empty? Need a signaling value here
+                        constant = make_constant(int, Shape([0]), [])
+                elif isinstance(arg, bool):
+                    constant = make_constant(OVType.boolean, Shape([]), [arg])
+                elif isinstance(arg, int):
+                    arg = maybe_convert_max_int(arg)
+                    constant = make_constant(OVType.i32, Shape([]), [arg])  # TODO: i32? why not i64?
+                elif isinstance(arg, float):
+                    constant = make_constant(OVType.f32, Shape([]), [arg])  # TODO: f32? why not f64?
+
+                if constant is None:
+                    if arg is None:
+                        self._inputs[i] = None
+                else:
+                    assert len(constant.outputs()) == 1
+                    result.append(constant.outputs()[0])
+                    self._inputs[i] = index
+                    index += 1
+        return result
+
+    def may_produce_alias(self, in_index: int, out_index: int) -> bool:
+        if self.get_op_type() in ["aten::conv1d", "aten::conv2d", "aten::conv3d"]:
+            # AliasDB::may_contain_alias sometimes return True for tensors produced by convnd, we have to workaround that
+            return False
+        try:
+            return self.alias_db.may_contain_alias(self._raw_input(in_index), self._raw_output(out_index))
+        except:
+            # Sometimes pytorch fails to get result with IndexError exception while these indexes exist in node
+            return False

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/backend.py

@@ -6,14 +6,20 @@
 
 import logging
 import os
+from functools import partial
+from hashlib import sha256
+
 import torch
 from torch._dynamo.backends.common import fake_tensor_unsupported
 from torch._dynamo.backends.registry import register_backend
 from torch._inductor.compile_fx import compile_fx
+from torch.fx.experimental.proxy_tensor import make_fx
 
 from openvino.frontend import FrontEndManager
 from openvino.runtime import Core, Type, PartialShape
-from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
+from openvino.frontend.pytorch.torchdynamo.execute import execute
 
 log = logging.getLogger(__name__)
 
@@ -36,6 +42,9 @@ log = logging.getLogger(__name__)
 @register_backend
 @fake_tensor_unsupported
 def openvino(subgraph, example_inputs):
+    if (os.getenv("PYTORCH_TRACING_MODE") is not None):
+        if (os.getenv("PYTORCH_TRACING_MODE") == "TORCHFX"):
+            return fx_openvino(subgraph, example_inputs)
     return ts_openvino(subgraph, example_inputs)
 
 
@@ -96,3 +105,25 @@ def ts_openvino(subgraph, example_inputs):
     except Exception as e:
         log.debug(f"Failed in compilation: {e}")
         return compile_fx(subgraph, example_inputs)
+
+
+def fx_openvino(subgraph, example_inputs):
+    try:
+        executor_parameters = None
+        if os.getenv("OPENVINO_TORCH_MODEL_CACHING") is not None:
+            model_hash_str = sha256(subgraph.code.encode('utf-8')).hexdigest()
+            executor_parameters = {"model_hash_str": model_hash_str}
+        model = make_fx(subgraph)(*example_inputs)
+        with torch.no_grad():
+            model.eval()
+        partitioner = Partitioner()
+        compiled_model = partitioner.make_partitions(model)
+
+        def _call(*args):
+            res = execute(compiled_model, *args, executor="openvino",
+                          executor_parameters=executor_parameters)
+            return res
+        return _call
+    except Exception as e:
+        log.debug(f"Failed in OpenVINO execution: {e}")
+        return compile_fx(subgraph, example_inputs)

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/compile.py

@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import os
+import torch
+import torch.overrides
+
+from torch.fx import GraphModule
+
+from openvino.frontend import FrontEndManager
+from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
+from openvino.runtime import Core, Type, PartialShape, serialize
+
+from typing import Callable, Optional
+
+
+def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None):
+    core = Core()
+
+    device = "CPU"
+
+    if os.getenv("OPENVINO_TORCH_BACKEND_DEVICE") is not None:
+        device = os.getenv("OPENVINO_TORCH_BACKEND_DEVICE")
+        assert device in core.available_devices, "Specified device " + device + " is not in the list of OpenVINO Available Devices"
+
+    file_name = None
+    cache_root = "./cache/"
+    if os.getenv("OPENVINO_TORCH_CACHE_DIR") is not None:
+        cache_root = os.getenv("OPENVINO_TORCH_CACHE_DIR")
+    if model_hash_str is not None:
+        model_cache_dir = cache_root + "/model/"
+        try:
+            os.makedirs(model_cache_dir, exist_ok=True)
+            file_name = model_cache_dir + model_hash_str + "_" + device
+        except OSError as error:
+            print("Cache directory ", cache_root, " cannot be created. Model caching is disabled. Error: ", error)
+            file_name = None
+            model_hash_str = None
+
+    if file_name is not None and os.path.isfile(file_name + ".xml") and os.path.isfile(file_name + ".bin"):
+        om = core.read_model(file_name + ".xml")
+    else:
+        fe_manager = FrontEndManager()
+        fe = fe_manager.load_by_framework("pytorch")
+
+        input_shapes = []
+        input_types = []
+        for idx, input_data in enumerate(args):  # subgraph.example_inputs):
+            input_types.append(input_data.type())
+            input_shapes.append(input_data.size())
+
+        decoder = TorchFXPythonDecoder(gm, gm, input_shapes=input_shapes, input_types=input_types)
+
+        im = fe.load(decoder)
+
+        om = fe.convert(im)
+
+        if file_name is not None:
+            serialize(om, file_name + ".xml", file_name + ".bin")
+
+    dtype_mapping = {
+        torch.float32: Type.f32,
+        torch.float64: Type.f64,
+        torch.float16: Type.f16,
+        torch.int64: Type.i64,
+        torch.int32: Type.i32,
+        torch.uint8: Type.u8,
+        torch.int8: Type.i8,
+        torch.bool: Type.boolean
+    }
+
+    for idx, input_data in enumerate(args):
+        om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
+        om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
+    om.validate_nodes_and_infer_types()
+
+    if model_hash_str is not None:
+        core.set_property({'CACHE_DIR': cache_root + '/blob'})
+
+    compiled = core.compile_model(om, device)
+    return compiled

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/execute.py

@@ -0,0 +1,151 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# mypy: ignore-errors
+
+from copy import deepcopy
+from dataclasses import dataclass
+from functools import lru_cache
+from types import MappingProxyType
+from warnings import warn
+
+import torch
+import torch.overrides
+
+from torch.fx import GraphModule
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
+
+from openvino.frontend import FrontEndManager
+from openvino.frontend.pytorch.fx_decoder import TorchFXPythonDecoder
+from openvino.frontend.pytorch.torchdynamo.partition import Partitioner
+from openvino.frontend.pytorch.torchdynamo.compile import openvino_compile
+from openvino.runtime import Core, Type, PartialShape
+
+from typing import Callable, Optional
+
+from torch.fx.experimental.proxy_tensor import make_fx, wrapper_and_args_for_make_fx
+
+
+DEFAULT_OPENVINO_PYTHON_CONFIG = MappingProxyType(
+    {
+        "use_python_fusion_cache": True,
+        "allow_single_op_fusion": True,
+    },
+)
+
+compiled_cache = {}
+max_openvino_partitions = 0
+partitioned_modules = {}
+
+
+def execute(
+    gm: GraphModule,
+    *args,
+    executor: str = "aten",
+    executor_parameters: Optional[dict] = None,
+):
+    if executor == "openvino":
+        return openvino_execute_partitioned(gm, *args, executor_parameters=executor_parameters)
+    elif executor == "strictly_openvino":
+        return openvino_execute(gm, *args, executor_parameters=executor_parameters)
+
+    msg = "Received unexpected value for 'executor': {0}. Allowed values are: openvino, strictly_openvino.".format(executor)
+    raise ValueError(msg)
+
+
+import numpy as np
+
+
+def openvino_execute(gm: GraphModule, *args, executor_parameters=None, partition_id):
+
+    executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
+
+    use_cache = executor_parameters.get(
+        "use_python_fusion_cache",
+        DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"],
+    )
+    global compiled_cache
+
+    model_hash_str = executor_parameters.get("model_hash_str", None)
+    if model_hash_str is not None:
+        model_hash_str = model_hash_str + str(partition_id)
+
+    if use_cache and (partition_id in compiled_cache):
+        compiled = compiled_cache[partition_id]
+    else:
+        compiled = openvino_compile(gm, *args, model_hash_str=model_hash_str)
+        compiled_cache[partition_id] = compiled
+
+    flat_args, _ = tree_flatten(args)
+    ov_inputs = [a.detach().cpu().numpy() for a in flat_args]
+
+    res = compiled(ov_inputs)
+
+    results1 = [torch.from_numpy(res[out]) for out in compiled.outputs]
+    if len(results1) == 1:
+        return results1[0]
+    return results1
+
+
+class OpenVINOGraphModule(torch.nn.Module):
+    def __init__(self, gm, partition_id, use_python_fusion_cache, model_hash_str: str = None):
+        super().__init__()
+        self.gm = gm
+        self.partition_id = partition_id
+        self.executor_parameters = {"use_python_fusion_cache": use_python_fusion_cache,
+                                    "model_hash_str": model_hash_str}
+        self.perm_fallback = False
+
+    def __call__(self, *args):
+        if self.perm_fallback:
+            return self.gm(*args)
+
+        try:
+            result = openvino_execute(self.gm, *args, executor_parameters=self.executor_parameters, partition_id=self.partition_id)
+        except Exception:
+            self.perm_fallback = True
+            return self.gm(*args)
+
+        return result
+
+
+def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_str: str = None):
+    global max_openvino_partitions
+    partition_id = max_openvino_partitions
+    for node in gm.graph.nodes:
+        # TODO: use a better way to identify fused submodule
+        if node.op == "call_module" and "fused_" in node.name:
+            openvino_submodule = getattr(gm, node.name)
+            gm.delete_submodule(node.target)
+            gm.add_submodule(
+                node.target,
+                OpenVINOGraphModule(openvino_submodule, partition_id, use_python_fusion_cache,
+                                    model_hash_str=model_hash_str),
+            )
+            partition_id = partition_id + 1
+
+    max_openvino_partitions = partition_id
+
+    return gm
+
+
+def openvino_execute_partitioned(gm: GraphModule, *args, executor_parameters=None):
+    executor_parameters = executor_parameters or DEFAULT_OPENVINO_PYTHON_CONFIG
+
+    global partitioned_modules
+
+    use_python_fusion_cache = executor_parameters.get(
+        "use_python_fusion_cache",
+        DEFAULT_OPENVINO_PYTHON_CONFIG["use_python_fusion_cache"],
+    )
+    model_hash_str = executor_parameters.get("model_hash_str", None)
+
+    signature = str(id(gm))
+    for idx, input_data in enumerate(args):
+        signature = signature + "_" + str(idx) + ":" + str(input_data.type())[6:] + ":" + str(input_data.size())[11:-1].replace(" ", "")
+
+    if signature not in partitioned_modules:
+        partitioned_modules[signature] = partition_graph(gm, use_python_fusion_cache=use_python_fusion_cache,
+                                                         model_hash_str=model_hash_str)
+    return partitioned_modules[signature](*args)

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py

@@ -0,0 +1,254 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+from typing import Dict
+
+import torch
+from torch.nn import Module
+from torch._ops import OpOverload
+
+from torch.fx import GraphModule
+from torch.fx.node import Node, _get_qualified_name
+from torch.fx.passes.operator_support import OperatorSupport
+from torch.fx.passes.tools_common import CALLABLE_NODE_OPS
+
+from torch.fx.experimental.proxy_tensor import DecompositionInterpreter
+from torch._decomp import decomposition_table
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
+
+import typing as t
+
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+def aten_to_dtype(self, dtype: torch.dtype, **kwargs):
+    if len(kwargs) > 0 or not dtype:
+        raise RuntimeError(
+            "No support for other to.dtype() formats other than to.dtype(self, dtype)"
+        )
+    return torch._prims.convert_element_type(self, dtype)
+
+
+# decomposition_table currently contains both aten2aten and aten2prim decomposition
+# this is a hack to separate them, as we only need aten2prim decomposition for nvfuser-supported aten graph lowering
+aten2aten_decomp = {}
+aten2prim_decomp = {}
+
+for op, decomp_fn in decomposition_table.items():
+    if "torch._refs" in decomp_fn.__module__:
+        aten2prim_decomp[op] = decomp_fn
+    else:
+        aten2aten_decomp[op] = decomp_fn
+
+aten2aten_decomp_skips = {
+    "aten.native_layer_norm_backward.default",
+    "aten.embedding_dense_backward.default",  # This is hurting nvfuser's perf
+    "aten.addmm.default",
+}
+
+for op, decomp_fn in decomposition_table.items():
+    if "torch._refs" in decomp_fn.__module__:
+        aten2prim_decomp[op] = decomp_fn
+    else:
+        if str(op) not in aten2aten_decomp_skips:
+            aten2aten_decomp[op] = decomp_fn
+
+
+aten2prim_decomp[torch.ops.aten.to.dtype] = aten_to_dtype
+
+
+class OperatorSupport(OperatorSupport):
+    """
+    Operator support for OpenVINO backend.
+    Currently, partitioning is based on FX ATen graph. The fused subgraph will latter be decomposed into prims.
+    To determine if an ATen ops is supported by nvFuser, we shall check the prim ops used in its ref decomposition.
+    Only if all the prim ops in the ref has a nvfuser_impl, we say this Aten op is suppported by nvFuser.
+    Note: When adding a rule, please add it to the corresponding section and follow the
+    alphabetical order.
+    """
+
+    def __init__(self):
+        # TODO: current list copied from torch/csrc/jit/codegen/cuda/parser.cpp is incorrect,
+        # as that file is solely for TorchScript and doesn't represent the actual status
+        # whether operation would be runnable by primTorch+nvFuser.
+        # We will iterate on this list to reflect the the reality.
+        """
+        support_dict = {
+            # ===============================================================
+            # call_function aten
+            # ===============================================================
+            # Following supported aten ops is copied from torch/csrc/jit/codegen/cuda/parser.cpp
+            # TODO: might need to update according to supported input types
+
+            "torch.ops.aten.relu": None,
+            "torch.ops.aten.relu_": None,
+            "torch.ops.aten.conv2d": None,
+            "torch.ops.aten._convolution": None,
+            "torch.ops.aten.convolution": None,
+            "torch.ops.aten.batch_norm": None,
+            "torch.ops.aten.layer_norm": None,
+            "torch.ops.aten.add": None,
+            "torch.ops.aten.add_": None,
+            "torch.ops.aten.mul": None,
+            "torch.ops.aten.mul_": None,
+            "torch.ops.aten.div": None,
+            "torch.ops.aten.floordiv": None,
+            "torch.ops.aten.tanh": None,
+            "torch.ops.aten.elu": None,
+            "torch.ops.aten.sigmoid": None,
+            "torch.ops.aten.gelu": None,
+            "torch.ops.aten.sqrt": None,
+            "torch.ops.aten.abs": None,
+            "torch.ops.aten.square": None,
+            "torch.ops.aten.hardtanh": None,
+            "torch.ops.aten.hardtanh_": None,
+            "torch.ops.aten.hardsigmoid": None,
+            "torch.ops.aten.hardswish": None,
+            "torch.ops.aten.hardswish_": None,
+            "torch.ops.aten.silu_": None,
+            "torch.ops.aten.relu6": None,
+            "torch.ops.aten.softmax": None,
+            "torch.ops.aten.matmul": None,
+            "torch.ops.aten.mm": None,
+            "torch.ops.aten.linear": None,
+            "torch.ops.aten.max_pool2d": None,
+            "torch.ops.aten.avg_pool2d": None,
+            "torch.ops.aten.adaptive_avg_pool2d": None,
+            "torch.ops.aten.adaptive_max_pool2d": None,
+            #"torch.ops.aten.max_pool2d_with_indices": None,
+            "torch.ops.aten.mean": None,
+            "torch.ops.aten.flatten": None,
+            #"torch.ops.prim.NumToTensor": None,
+            "torch.ops.aten.contiguous": None,
+            "torch.ops.aten.as_tensor": None,
+            "torch.ops.aten.Int": None,
+            "torch.ops.aten.to": None,
+            "torch.ops.aten.permute": None,
+            "torch.ops.aten.embedding": None,
+            "torch.ops.aten.transpose": None,
+            "torch.ops.aten.size": None,
+            "torch.ops.aten.view": None,
+            "torch.ops.aten.unsqueeze": None,
+            "torch.ops.aten.rsub": None,
+            "torch.ops.aten.slice": None,
+            #"torch.ops.prim.Loop": None,
+            #"torch.ops.prim.If": None,
+            #"torch.ops.prim.Constant": None,
+            "torch.ops.aten.dim": None,
+            "torch.ops.aten.reciprocal": None,
+            "torch.ops.aten.sub": None,
+            "torch.ops.aten.eq": None,
+            "torch.ops.aten.ne": None,
+            "torch.ops.aten.gt": None,
+            "torch.ops.aten.lt": None,
+            "torch.ops.aten.neg": None,
+            #"torch.ops.prim.TupleConstruct": None,
+            "torch.ops.aten.append": None,
+            "getattr": None,
+            "_operator.getitem": None,
+        }
+        """
+        # Just added Resnet50 supported iterations
+        support_dict = {
+            "_operator.getitem": None,
+            "torch.ops.aten._adaptive_avg_pool2d.default": None,
+            "torch.ops.aten._softmax.default": None,
+            "torch.ops.aten._to_copy.default": None,
+            "torch.ops.aten._unsafe_view.default": None,
+            "torch.ops.aten._unsafe_view.default": None,
+            "torch.ops.aten.add.Tensor": None,
+            "torch.ops.aten.add_.Tensor": None,
+            "torch.ops.aten.addmm.default": None,
+            "torch.ops.aten.arange.start": None,
+            "torch.ops.aten.arange.default": None,
+            "torch.ops.aten.argmax.default": None,
+            "torch.ops.aten.avg_pool2d.default": None,
+            "torch.ops.aten.bitwise_and.Tensor": None,
+            "torch.ops.aten.bmm.default": None,
+            "torch.ops.aten.cat.default": None,
+            "torch.ops.aten.clamp_min.default": None,
+            "torch.ops.aten.clone.default": None,
+            "torch.ops.aten.convolution.default": None,
+            "torch.ops.aten.copy_.default": None,
+            "torch.ops.aten.cos.default": None,
+            "torch.ops.aten.cumsum.default": None,
+            "torch.ops.aten.detach.default": None,
+            "torch.ops.aten.div.Scalar": None,
+            "torch.ops.aten.div.Tensor": None,
+            "torch.ops.aten.embedding.default": None,
+            "torch.ops.aten.empty.memory_format": None,
+            "torch.ops.aten.eq.Scalar": None,
+            "torch.ops.aten.eq.Tensor": None,
+            "torch.ops.aten.exp.default": None,
+            "torch.ops.aten.expand.default": None,
+            "torch.ops.aten.full.default": None,
+            "torch.ops.aten.gather.default": None,
+            "torch.ops.aten.gelu.default": None,
+            "torch.ops.aten.gt.Scalar": None,
+            "torch.ops.aten.hardsigmoid.default": None,
+            "torch.ops.aten.hardswish_.default": None,
+            "torch.ops.aten.hardtanh_.default": None,
+            "torch.ops.aten.index.Tensor": None,
+            "torch.ops.aten.lift_fresh_copy.default": None,
+            "torch.ops.aten.linalg_vector_norm.default": None,
+            "torch.ops.aten.lt.Tensor": None,
+            "torch.ops.aten.log.default": None,
+            "torch.ops.aten.logsumexp.default": None,
+            "torch.ops.aten.masked_fill_.Scalar": None,
+            "torch.ops.aten.masked_fill.Tensor": None,
+            "torch.ops.aten.max_pool2d_with_indices.default": None,
+            "torch.ops.aten.mean.dim": None,
+            "torch.ops.aten.mm.default": None,
+            "torch.ops.aten.mul.Scalar": None,
+            "torch.ops.aten.mul.Tensor": None,
+            "torch.ops.aten.native_batch_norm.default": None,
+            "torch.ops.aten.native_group_norm.default": None,
+            "torch.ops.aten.native_layer_norm.default": None,
+            "torch.ops.aten.neg.default": None,
+            "torch.ops.aten.new_ones.default": None,
+            "torch.ops.aten.permute.default": None,
+            "torch.ops.aten.pow.Tensor_Scalar": None,
+            "torch.ops.aten.relu.default": None,
+            "torch.ops.aten.relu_.default": None,
+            "torch.ops.aten.rsub.Scalar": None,
+            "torch.ops.aten.select.int": None,
+            "torch.ops.aten.sigmoid.default": None,
+            "torch.ops.aten.silu.default": None,
+            "torch.ops.aten.silu_.default": None,
+            "torch.ops.aten.sin.default": None,
+            "torch.ops.aten.slice.Tensor": None,
+            "torch.ops.aten.split.Tensor": None,
+            "torch.ops.aten.sub.default": None,
+            "torch.ops.aten.sub.Tensor": None,
+            "torch.ops.aten.t.default": None,
+            "torch.ops.aten.tanh.default": None,
+            "torch.ops.aten.transpose.int": None,
+            "torch.ops.aten.unsqueeze.default": None,
+            "torch.ops.aten.upsample_nearest2d.default": None,
+            "torch.ops.aten.view.default": None,
+            "torch.ops.aten.where.self": None,
+            "torch.ops.aten.zeros_like.default": None,
+        }
+
+        super().__init__(support_dict)
+
+    def is_node_supported(self, submodules: t.Mapping[str, Module], node: Node) -> bool:
+        # OpenVINO FX subgraph should be purely functional
+        if node.op not in CALLABLE_NODE_OPS:
+            return False
+
+        # ops in supported_dict doesn't have overload name
+        # use overloadpacket's qualified_name for OpOverload
+        if isinstance(node.target, OpOverload):
+            target = _get_qualified_name(node.target.overloadpacket)
+
+            if target in self._support_dict:
+                return True
+
+        return super().is_node_supported(submodules, node)

src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/partition.py

@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# mypy: ignore-errors
+
+from typing import Dict
+
+import torch
+from torch.nn import Module
+from torch.fx import GraphModule, Node
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
+
+from torch.fx.experimental.proxy_tensor import DecompositionInterpreter
+from torch._decomp import decomposition_table
+from torch.fx.experimental.proxy_tensor import make_fx
+from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten
+from openvino.frontend.pytorch.torchdynamo.op_support import OperatorSupport
+
+import typing as t
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class Partitioner:
+    def __init__(self):
+        self.supported_ops = OperatorSupport()
+
+    def fx_serialize(self, graph_module: GraphModule, *args, **kwargs):
+        fx_gm = make_fx(graph_module)(*args)
+        return fx_gm  # prim_module
+
+    def add_get_attr_inputs(self, partitions: t.List[Partition]):
+        # TODO: Find a more efficient way to include input
+        # "get_attr" nodes to the partitions.
+        getattr_to_merge: Dict[Node, Node] = {}
+        for partition in partitions:
+            for pnode in partition.nodes:
+                for pnode_input in pnode.all_input_nodes:
+                    if pnode_input.op in ["get_attr"] and pnode_input.op not in getattr_to_merge:
+                        getattr_to_merge[pnode_input] = partition
+        for getattr_node, getattr_part in getattr_to_merge.items():
+            getattr_part.add_node(getattr_node)
+
+    def make_partitions(self, graph_module: GraphModule) -> GraphModule:
+        # entry function for nvFuser backend
+        # FX graph based partitioning based on nvfuser supported ops
+        partitioner = CapabilityBasedPartitioner(
+            graph_module, self.supported_ops, allows_single_node_partition=False)
+        partitions = partitioner.propose_partitions()
+        self.add_get_attr_inputs(partitions)
+        fused_graph_module = partitioner.fuse_partitions(partitions)
+
+        return fused_graph_module

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

@@ -7,6 +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, Tensor
+from openvino.frontend.pytorch.utils import ivalue_to_constant, get_value_from_getattr, pt_to_ov_type_map
 from openvino.runtime import opset11 as ops
 
 import typing
@@ -22,99 +23,11 @@ class ModelWrapper(torch.nn.Module):
     def __init__(self, model):
         super().__init__()
         self.model = model
-    
+
     def forward(self, {input_sign}):
         return self.model({example_input})
 """
 
-
-def get_type_from_py_type(value):
-    if isinstance(value, float):
-        return OVType.f32
-    if isinstance(value, bool):
-        return OVType.boolean
-    if isinstance(value, int):
-        # Python int is 64 bit, but we will convert it to int32 except cases when it can't fit in 32 bits
-        if torch.iinfo(torch.int).min <= value <= torch.iinfo(torch.int).max:
-            return OVType.i32
-        return OVType.i64
-    return OVType.dynamic
-
-
-def ivalue_to_constant(ivalue):
-    ov_type = get_type_from_py_type(ivalue)
-    if ov_type.is_static():
-        return op.Constant(ov_type, Shape([]), [ivalue]).outputs()
-
-    if isinstance(ivalue, (list, tuple)):
-        assert len(ivalue) > 0, "Can't deduce type for empty list"
-        ov_type = get_type_from_py_type(ivalue[0])
-        assert ov_type.is_static(), "Can't deduce type for list"
-        return op.Constant(ov_type, Shape([len(ivalue)]), ivalue).outputs()
-
-    if isinstance(ivalue, torch.Tensor):
-        ivalue = ivalue.to(memory_format=torch.contiguous_format)
-        if ivalue.dtype == torch.bfloat16:
-            # reinterpret bfloat16 data as float16 to allow conversion to numpy
-            ivalue = ivalue.view(torch.float16)
-            narr = ivalue.numpy(force=True)
-            if not narr.flags['C_CONTIGUOUS']:
-                narr = np.ascontiguousarray(narr)
-            # TODO: this tensor doesn't share memory with initial tensor
-            tensor = Tensor(narr, ivalue.shape, OVType.bf16)
-            ov_const = op.Constant(tensor, shared_memory=True)
-        else:
-            narr = ivalue.numpy(force=True)
-            if not narr.flags['C_CONTIGUOUS']:
-                narr = np.ascontiguousarray(narr)
-            ov_const = op.Constant(narr, shared_memory=True)
-        return ov_const.outputs()
-    return None
-
-
-def get_value_from_getattr(getattr_node, self_module):
-    assert getattr_node.kind() == "prim::GetAttr", "Got node of kind not equal to prim::GetAttr"
-    # GetAttr nodes can be nested
-    stack = []
-    while getattr_node.kind() == "prim::GetAttr":
-        stack.append(getattr_node)
-        inputs = list(getattr_node.inputs())
-        if len(inputs) == 0:
-            break
-        getattr_node = inputs[0].node()
-    module = self_module
-    while len(stack) > 0:
-        node = stack.pop()
-        attr_name = node.s("name")
-        assert hasattr(module, attr_name), f"No attribute with name \"{attr_name}\" found in module."
-        module = getattr(module, attr_name)
-    return module
-
-
-pt_to_ov_type_map = {
-    "float": OVType.f32,
-    "int": OVType.i32,
-    "bool": OVType.boolean,
-    "torch.bfloat16": OVType.bf16,
-    "torch.float16": OVType.f16,
-    "torch.float32": OVType.f32,
-    "torch.float64": OVType.f64,
-    "torch.uint8": OVType.u8,
-    "torch.int8": OVType.i8,
-    "torch.int32": OVType.i32,
-    "torch.int64": OVType.i64,
-    "torch.bool": OVType.boolean,
-    "torch.DoubleTensor": OVType.f64,
-    "torch.FloatTensor": OVType.f32,
-    "torch.IntTensor": OVType.i32,
-    "torch.LongTensor": OVType.i64,
-    "torch.BoolTensor": OVType.boolean,
-    "torch.quint8": OVType.u8,
-    "torch.qint8": OVType.i8,
-    "torch.qint32": OVType.i32
-}
-
-
 class TorchScriptPythonDecoder (Decoder):
     def __init__(self, pt_module, graph_element=None, example_input=None, alias_db=None):
         Decoder.__init__(self)
@@ -171,7 +84,7 @@ class TorchScriptPythonDecoder (Decoder):
 
             # PyTorch has some difficulties to trace models with named unordered parameters:
             # torch < 2.0.0 supports only positional arguments for tracing
-            # pytorch == 2.0.0 supports input kwargs tracing, 
+            # pytorch == 2.0.0 supports input kwargs tracing,
             # but does not support complex nested objects (e. g. tuple of tuples of tensors)
             # We will use wrapper for making them positional as workaround.
 
@@ -535,6 +448,9 @@ class TorchScriptPythonDecoder (Decoder):
             # Sometimes pytorch fails to get result with IndexError exception while these indexes exist in node
             return False
 
+    def inlined_inputs(self, index):
+        return []
+
     @staticmethod
     def _transform_tensor_list_constants_to_listconstruct(graph: torch.Graph):
         # Function replaces prim::Constant containing List of Tensors with
@@ -578,4 +494,4 @@ class TorchScriptPythonDecoder (Decoder):
             const_input = graph.insertConstant(value)
             const_input.node().moveBefore(node)
             const_input.node().copyMetadata(node)
-            node.output().replaceAllUsesWith(const_input)
+            node.output().replaceAllUsesWith(const_input)

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

@@ -0,0 +1,137 @@
+
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+# flake8: noqa
+# mypy: ignore-errors
+
+import typing
+from packaging.version import parse
+import torch
+import numpy as np
+import inspect
+import ctypes
+
+from openvino.runtime import op, PartialShape, Type as OVType, OVAny, Shape, Tensor
+
+
+def maybe_convert_max_int(value : int):
+    # FIXME: This is a convertion from 64-bit positive max integer value
+    # to 32-bit positive max integer value. Find a better way to handle this.
+    if value == torch.iinfo(torch.int64).max:
+        return torch.iinfo(torch.int32).max
+    else:
+        return value
+
+def make_constant(*args, **kwargs):
+    return op.Constant(*args, **kwargs)
+
+def fetch_attr(self_module, target : str):
+    """
+    Fetch an attribute from the ``Module`` hierarchy of ``self.module``.
+
+    Args:
+        target (str): The fully-qualified name of the attribute to fetch
+
+    Return:
+        Any: The value of the attribute.
+    """
+    target_atoms = target.split('.')
+    attr_itr = self_module
+    for i, atom in enumerate(target_atoms):
+        if not hasattr(attr_itr, atom):
+            raise RuntimeError(f"Node referenced nonexistent target {'.'.join(target_atoms[:i])}")
+        attr_itr = getattr(attr_itr, atom)
+    return attr_itr
+
+
+def get_type_from_py_type(value):
+    if isinstance(value, float):
+        return OVType.f32
+    if isinstance(value, bool):
+        return OVType.boolean
+    if isinstance(value, int):
+        # Python int is 64 bit, but we will convert it to int32 except cases when it can't fit in 32 bits
+        if torch.iinfo(torch.int).min <= value <= torch.iinfo(torch.int).max:
+            return OVType.i32
+        return OVType.i64
+    return OVType.dynamic
+
+
+def ivalue_to_constant(ivalue):
+    ov_type = get_type_from_py_type(ivalue)
+    if ov_type.is_static():
+        return op.Constant(ov_type, Shape([]), [ivalue]).outputs()
+
+    if isinstance(ivalue, (list, tuple)):
+        assert len(ivalue) > 0, "Can't deduce type for empty list"
+        ov_type = get_type_from_py_type(ivalue[0])
+        assert ov_type.is_static(), "Can't deduce type for list"
+        return op.Constant(ov_type, Shape([len(ivalue)]), ivalue).outputs()
+
+    if isinstance(ivalue, torch.Tensor):
+        ivalue = ivalue.to(memory_format=torch.contiguous_format)
+        if ivalue.dtype == torch.bfloat16:
+            # reinterpret bfloat16 data as float16 to allow conversion to numpy
+            ivalue = ivalue.view(torch.float16)
+            narr = ivalue.numpy(force=True)
+            if not narr.flags['C_CONTIGUOUS']:
+                narr = np.ascontiguousarray(narr)
+            # TODO: this tensor doesn't share memory with initial tensor
+            tensor = Tensor(narr, ivalue.shape, OVType.bf16)
+            ov_const = op.Constant(tensor, shared_memory=True)
+        else:
+            narr = ivalue.numpy(force=True)
+            if not narr.flags['C_CONTIGUOUS']:
+                narr = np.ascontiguousarray(narr)
+            ov_const = op.Constant(narr, shared_memory=True)
+        return ov_const.outputs()
+    return None
+
+def get_value_from_getattr(getattr_node, self_module):
+    assert getattr_node.kind() == "prim::GetAttr", "Got node of kind not equal to prim::GetAttr"
+    # GetAttr nodes can be nested
+    stack = []
+    while getattr_node.kind() == "prim::GetAttr":
+        stack.append(getattr_node)
+        inputs = list(getattr_node.inputs())
+        if len(inputs) == 0:
+            break
+        getattr_node = inputs[0].node()
+    module = self_module
+    while len(stack) > 0:
+        node = stack.pop()
+        attr_name = node.s("name")
+        assert hasattr(module, attr_name), f"No attribute with name \"{attr_name}\" found in module."
+        module = getattr(module, attr_name)
+    return module
+
+pt_to_ov_type_map = {
+    "float": OVType.f32,
+    "int": OVType.i32,
+    "bool": OVType.boolean,
+    "torch.bfloat16": OVType.bf16,
+    "torch.float16": OVType.f16,
+    "torch.float32": OVType.f32,
+    "torch.float64": OVType.f64,
+    "torch.uint8": OVType.u8,
+    "torch.int8": OVType.i8,
+    "torch.int32": OVType.i32,
+    "torch.int64": OVType.i64,
+    "torch.bool": OVType.boolean,
+    "torch.DoubleTensor": OVType.f64,
+    "torch.FloatTensor": OVType.f32,
+    "torch.IntTensor": OVType.i32,
+    "torch.LongTensor": OVType.i64,
+    "torch.BoolTensor": OVType.boolean,
+    "torch.quint8": OVType.u8,
+    "torch.qint8": OVType.i8,
+    "torch.qint32": OVType.i32
+}
+
+ov_to_c_type_map = {
+    OVType.f32: ctypes.c_float,
+    OVType.f64: ctypes.c_double,
+    OVType.i32: ctypes.c_int,
+    OVType.i64: ctypes.c_int64,
+}

src/bindings/python/src/pyopenvino/frontend/pytorch/decoder.hpp

@@ -113,6 +113,9 @@ class PyDecoder : public ov::frontend::pytorch::TorchDecoder {
     bool may_produce_alias(size_t in_index, size_t out_index) const override {
         PYBIND11_OVERRIDE_PURE(bool, TorchDecoder, may_produce_alias, in_index, out_index);
     }
+
+    ov::OutputVector inlined_inputs(size_t start_index) const override {
+        PYBIND11_OVERRIDE_PURE(ov::OutputVector, TorchDecoder, inlined_inputs, start_index); }
 };
 
 void regclass_frontend_pytorch_decoder(py::module m);

src/frontends/pytorch/include/openvino/frontend/pytorch/decoder.hpp

@@ -109,6 +109,12 @@ public:
 
     /// \brief Returns if output may contain alias of input in AliasDB
     virtual bool may_produce_alias(size_t in_index, size_t out_index) const = 0;
+
+    /// Returns new nodes for inputs inlined in the op itself
+    // Used in Torch.FX decoder
+    virtual OutputVector inlined_inputs(size_t start_index) const {
+        return {};
+    }
 };
 
 }  // namespace pytorch

src/frontends/pytorch/src/frontend.cpp

@@ -38,6 +38,7 @@
 #include "transforms/prim_tuple_unpack_parameter_replacer.hpp"
 #include "transforms/quantized_node_remover.hpp"
 #include "transforms/rfftn_complex_replacer.hpp"
+#include "transforms/softmax_reshape_elimination.hpp"
 #include "transforms/string_equality_replacer.hpp"
 #include "transforms/tuple_unpack_replacer.hpp"
 #include "translate_session.hpp"
@@ -109,7 +110,14 @@ std::string pack_detailed_failure_report(const std::map<std::string, std::string
 }
 }  // namespace
 
-FrontEnd::FrontEnd() : m_op_translators(get_supported_ops()) {}
+FrontEnd::FrontEnd() {
+    const char* torch_tracing_mode = std::getenv("PYTORCH_TRACING_MODE");
+    if ((torch_tracing_mode != nullptr) && std::strcmp(torch_tracing_mode, "TORCHFX") == 0) {
+        m_op_translators = get_supported_ops_fx();
+    } else {
+        m_op_translators = get_supported_ops_ts();
+    }
+}
 
 std::shared_ptr<Model> FrontEnd::convert(const InputModel::Ptr& model) const {
     auto converted_model = convert_partially(model);
@@ -186,6 +194,7 @@ void FrontEnd::normalize(const std::shared_ptr<ov::Model>& model) const {
     manager.register_pass<ov::frontend::pytorch::pass::IndexLoopGetitemReplacer>();
     manager.register_pass<ov::frontend::pytorch::pass::DequantizeNodeRemover>();
     manager.register_pass<ov::frontend::pytorch::pass::QuantizedNodeRemover>();
+    manager.register_pass<ov::frontend::pytorch::pass::SoftmaxReshapeElimination>();
     manager.register_pass<ov::pass::RemoveMultiSubGraphOpDanglingParamsResults>();
     manager.register_pass<ov::pass::ReverseShapeAndTypeInfer>();
 

src/frontends/pytorch/src/op/addmm.cpp

@@ -17,14 +17,21 @@ namespace op {
 using namespace ov::op;
 
 OutputVector translate_addmm(const NodeContext& context) {
-    num_inputs_check(context, 5, 5);
+    num_inputs_check(context, 3, 5);
     auto input = context.get_input(0);
     auto m1 = context.get_input(1);
     auto m2 = context.get_input(2);
-    auto beta = context.get_input(3);
-    auto alpha = context.get_input(4);
-    auto beta_converted = context.mark_node(std::make_shared<v1::ConvertLike>(beta, input));
     auto mm = context.mark_node(std::make_shared<v0::MatMul>(m1, m2));
+    auto one = context.mark_node(v0::Constant::create(element::f32, Shape{}, {1}));
+    ov::Output<Node> alpha = one;
+    ov::Output<Node> beta = one;
+    if (!context.input_is_none(3)) {
+        beta = context.get_input(3);
+    }
+    if (!context.input_is_none(4)) {
+        alpha = context.get_input(4);
+    }
+    auto beta_converted = context.mark_node(std::make_shared<v1::ConvertLike>(beta, input));
     auto alpha_converted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, mm));
     auto input_beta = context.mark_node(std::make_shared<v1::Multiply>(input, beta_converted));
     auto mm_alpha = context.mark_node(std::make_shared<v1::Multiply>(mm, alpha_converted));

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

@@ -85,7 +85,80 @@ OutputVector translate_arange(const NodeContext& context) {
     return {range};
 };
 
+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)
+        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 {
+        FRONT_END_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 {
+            FRONT_END_OP_CONVERSION_CHECK(false, "Couldn't get dtype input");
+        }
+    }
+    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));
+    }
+    return {range};
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

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, 3, 3);
+    num_inputs_check(context, 2, 3);
     auto input = context.get_input(0);
     bool keep_dims = false;
     auto k = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1}));
@@ -69,4 +69,4 @@ OutputVector translate_argmin(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

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, 6, 7);
+    num_inputs_check(context, 3, 7);
     auto input = context.get_input(0);
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
@@ -30,9 +30,21 @@ OutputVector translate_avg_poolnd(const NodeContext& context) {
         // In case strides are not provided default is kernel
         strides = kernel;
     }
-    auto pads = context.const_input<Shape>(3);  // pytorch supports only symmetric padding
-    auto rounding_type = context.const_input<bool>(4) ? ov::op::RoundingType::CEIL : ov::op::RoundingType::FLOOR;
-    auto count_include_pad = context.const_input<bool>(5);
+    Shape pads;
+    bool count_include_pad = true;
+    if (context.input_is_none(3)) {
+        count_include_pad = false;
+        pads = Shape(kernel.size(), 0);
+    } else {
+        pads = context.const_input<Shape>(3);  // pytorch supports only symmetric padding
+    }
+    ov::op::RoundingType rounding_type = ov::op::RoundingType::FLOOR;
+    if (!(context.input_is_none(4))) {
+        rounding_type = context.const_input<bool>(4) ? ov::op::RoundingType::CEIL : ov::op::RoundingType::FLOOR;
+    }
+    if (!(context.input_is_none(5))) {
+        count_include_pad = context.const_input<bool>(5);
+    }
     FRONT_END_OP_CONVERSION_CHECK(context.input_is_none(6),
                                   "Translation for aten::avg_pool2d do not support divisor_override input.");
     // Although ov::AvgPool provides exclude_pad=false,

src/frontends/pytorch/src/op/batch_norm.cpp

@@ -16,6 +16,7 @@
 #include "openvino/op/squeeze.hpp"
 #include "openvino/op/subtract.hpp"
 #include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -97,7 +98,12 @@ OutputVector translate_batch_norm(const NodeContext& context) {
         std::make_shared<v5::BatchNormInference>(input, weight, bias, running_mean, running_var, epsilon))};
 };
 
+OutputVector translate_batch_norm_fx(const NodeContext& context) {
+    auto output = translate_batch_norm(context);
+    return {context.mark_node(make_list_construct(output))};
+}
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/bitwise.cpp

@@ -0,0 +1,51 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/logical_and.hpp"
+#include "openvino/op/logical_not.hpp"
+#include "openvino/op/logical_or.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_bitwise_not(const NodeContext& context) {
+    num_inputs_check(context, 1, 2);
+    auto x = context.get_input(0);
+    FRONT_END_OP_CONVERSION_CHECK(x.get_element_type().compatible(element::boolean),
+                                  "aten::bitwise_not suppored only for boolean input");
+    auto not_x = context.mark_node(std::make_shared<ov::op::v1::LogicalNot>(x));
+    if (!context.input_is_none(1)) {
+        context.mutate_input(1, not_x);
+    }
+    return {not_x};
+};
+
+OutputVector translate_bitwise_and(const NodeContext& context) {
+    num_inputs_check(context, 2, 2);
+    auto x = context.get_input(0);
+    auto y = context.get_input(1);
+    FRONT_END_OP_CONVERSION_CHECK(x.get_element_type().compatible(element::boolean),
+                                  "aten::bitwise_not suppored only for boolean input");
+    auto and_x = context.mark_node(std::make_shared<ov::op::v1::LogicalAnd>(x, y));
+    return {and_x};
+};
+
+OutputVector translate_bitwise_or(const NodeContext& context) {
+    num_inputs_check(context, 2, 2);
+    auto x = context.get_input(0);
+    auto y = context.get_input(1);
+    FRONT_END_OP_CONVERSION_CHECK(x.get_element_type().compatible(element::boolean),
+                                  "aten::bitwise_not suppored only for boolean input");
+    auto or_x = context.mark_node(std::make_shared<ov::op::v1::LogicalOr>(x, y));
+    return {or_x};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op/bitwise_not.cpp

@@ -1,29 +0,0 @@
-// Copyright (C) 2018-2023 Intel Corporation
-// SPDX-License-Identifier: Apache-2.0
-//
-
-#include "openvino/frontend/pytorch/node_context.hpp"
-#include "openvino/op/logical_not.hpp"
-#include "utils.hpp"
-
-namespace ov {
-namespace frontend {
-namespace pytorch {
-namespace op {
-
-OutputVector translate_bitwise_not(const NodeContext& context) {
-    num_inputs_check(context, 1, 2);
-    auto x = context.get_input(0);
-    FRONT_END_OP_CONVERSION_CHECK(x.get_element_type().compatible(element::boolean),
-                                  "aten::bitwise_not suppored only for boolean input");
-    auto not_x = context.mark_node(std::make_shared<ov::op::v1::LogicalNot>(x));
-    if (!context.input_is_none(1)) {
-        context.mutate_input(1, not_x);
-    }
-    return {not_x};
-};
-
-}  // namespace op
-}  // namespace pytorch
-}  // namespace frontend
-}  // namespace ov

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

@@ -15,11 +15,9 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_cat(const NodeContext& context) {
-    // This translator is only needed to get axis as constant from external scope
-    num_inputs_check(context, 2, 2);
-    const auto&& list_elems = get_list_as_outputs(context.get_input(0));
-    auto axis = context.const_input<int64_t>(1);
+OutputVector translate_cat_common(const NodeContext& context,
+                                  const std::deque<ov::Output<ov::Node>>& list_elems,
+                                  int64_t axis) {
     if (list_elems.empty()) {
         // couldn't get list elements
         auto fw_node = std::make_shared<PtFrameworkNode>(context.get_decoder(), OutputVector{context.get_input(0)}, 1);
@@ -36,9 +34,28 @@ OutputVector translate_cat(const NodeContext& context) {
     }
     auto concat = std::make_shared<v0::Concat>(OutputVector(list_elems.begin(), list_elems.end()), axis);
     return {context.mark_node(concat)};
+}
+
+OutputVector translate_cat(const NodeContext& context) {
+    // This translator is only needed to get axis as constant from external scope
+    num_inputs_check(context, 2, 2);
+    const auto&& list_elems = get_list_as_outputs(context.get_input(0));
+    auto axis = context.const_input<int64_t>(1);
+    return translate_cat_common(context, list_elems, axis);
+};
+
+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());
+    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);
+    return translate_cat_common(context, list_elems, axis);
 };
 
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/div.cpp

@@ -50,4 +50,4 @@ OutputVector translate_div(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/embedding.cpp

@@ -16,7 +16,7 @@ namespace op {
 OutputVector translate_embedding(const NodeContext& context) {
     // aten::embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool
     // sparse=False)
-    num_inputs_check(context, 5, 5);
+    num_inputs_check(context, 2, 5);
     auto data = context.get_input(0);
     auto indices = context.get_input(1);
     indices = context.mark_node(std::make_shared<ov::op::v0::Convert>(indices, element::i32));
@@ -32,4 +32,4 @@ OutputVector translate_embedding(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/expand.cpp

@@ -49,6 +49,24 @@ OutputVector translate_expand_as(const NodeContext& context) {
     return base_expand(context, x, sizes);
 };
 
+OutputVector translate_expand_fx(const NodeContext& context) {
+    // aten::expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a)
+    num_inputs_check(context, 2, 3);
+    auto x = context.get_input(0);
+    // TODO: This is a temporary solution to optimize out Broadcast if the input and
+    // output shapes are same. This should be removed after a proper optimization is
+    // implemented.
+    auto sizes_const = context.const_input<Shape>(1);
+    if (x.get_shape() == sizes_const) {
+        return {x};
+    }
+    auto sizes = context.get_input(1);
+    // TODO: figure out what implicit means
+    FRONT_END_OP_CONVERSION_CHECK(context.input_is_none(2) || context.const_input<bool>(2) == false,
+                                  "Unexpected value of implicit for expand operation");
+    return base_expand(context, x, sizes);
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/op/gather.cpp

@@ -15,7 +15,7 @@ namespace op {
 OutputVector translate_gather(const NodeContext& context) {
     // aten::gather(Tensor self, int dim, Tensor index, *, bool sparse_grad=False) -> Tensor
     // aten::gather.out(Tensor self, int dim, Tensor index, *, bool sparse_grad=False, Tensor(a!) out) -> Tensor(a!)
-    num_inputs_check(context, 4, 5);
+    num_inputs_check(context, 3, 5);
     auto x = context.get_input(0);
     auto axis = context.const_input<int64_t>(1);
     auto index = context.get_input(2);
@@ -31,4 +31,4 @@ OutputVector translate_gather(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/gelu.cpp

@@ -13,9 +13,12 @@ namespace pytorch {
 namespace op {
 
 OutputVector translate_gelu(const NodeContext& context) {
-    num_inputs_check(context, 2, 2);
+    num_inputs_check(context, 1, 2);
     auto x = context.get_input(0);
-    auto approximate = context.const_input<std::string>(1);
+    std::string approximate = "none";
+    if (!context.input_is_none(1)) {
+        approximate = context.const_input<std::string>(1);
+    }
     if (approximate == "none") {
         return {context.mark_node(std::make_shared<ov::op::v7::Gelu>(x, ov::op::GeluApproximationMode::ERF))};
     }
@@ -28,4 +31,4 @@ OutputVector translate_gelu(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/group_norm.cpp

@@ -11,6 +11,7 @@
 #include "openvino/op/reshape.hpp"
 #include "openvino/op/subtract.hpp"
 #include "openvino/op/unsqueeze.hpp"
+#include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -20,14 +21,19 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_group_norm(const NodeContext& context) {
+OutputVector translate_group_norm_common(const NodeContext& context,
+                                         const size_t max_inputs,
+                                         const size_t group_idx,
+                                         const size_t eps_idx,
+                                         const size_t weights_idx,
+                                         const size_t bias_idx) {
     // aten::group_norm(Tensor input, int num_groups, Tensor? weight=None, Tensor? bias=None, float
     // eps=1.0000000000000001e-05, bool cudnn_enabled=True) -> Tensor
-    num_inputs_check(context, 2, 6);
+    num_inputs_check(context, 2, max_inputs);
     auto data = context.get_input(0);
-    auto num_groups = context.const_input<int64_t>(1);
+    auto num_groups = context.const_input<int64_t>(group_idx);
     // input 2 - weights and input 3 - bias are optional without default value, we handle them later
-    auto eps = static_cast<float>(context.const_input<double>(4));
+    auto eps = static_cast<float>(context.const_input<double>(eps_idx));
     Output<Node> input_shape;
     Output<Node> input_rank;
     std::tie(input_shape, input_rank) = get_shape_rank(context, data, true, element::i32);
@@ -41,13 +47,13 @@ OutputVector translate_group_norm(const NodeContext& context) {
     auto norm = context.mark_node(std::make_shared<v1::Reshape>(reshaped_norm, input_shape, true));
     auto skip_last = context.mark_node(std::make_shared<v1::Subtract>(input_rank, scalar_one));
     auto axes = context.mark_node(std::make_shared<v4::Range>(scalar_one, skip_last, scalar_one, element::i32));
-    if (!context.input_is_none(2)) {
-        auto weights = context.get_input(2);
+    if (!context.input_is_none(weights_idx)) {
+        auto weights = context.get_input(static_cast<int>(weights_idx));
         weights = context.mark_node(std::make_shared<v0::Unsqueeze>(weights, axes));
         norm = context.mark_node(std::make_shared<v1::Multiply>(norm, weights));
     }
-    if (!context.input_is_none(3)) {
-        auto bias = context.get_input(3);
+    if (!context.input_is_none(bias_idx)) {
+        auto bias = context.get_input(static_cast<int>(bias_idx));
         bias = context.mark_node(std::make_shared<v0::Unsqueeze>(bias, axes));
         norm = context.mark_node(std::make_shared<v1::Add>(norm, bias));
     }
@@ -55,7 +61,16 @@ OutputVector translate_group_norm(const NodeContext& context) {
     return {norm};
 };
 
+OutputVector translate_group_norm(const NodeContext& context) {
+    return translate_group_norm_common(context, 6, 1, 4, 2, 3);
+}
+
+OutputVector translate_group_norm_fx(const NodeContext& context) {
+    auto output = translate_group_norm_common(context, 8, 6, 7, 1, 2);
+    return {context.mark_node(make_list_construct(output))};
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/index.cpp

@@ -0,0 +1,38 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/concat.hpp"
+#include "openvino/op/convert.hpp"
+#include "openvino/op/gather_nd.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+OutputVector translate_index_fx(const NodeContext& context) {
+    num_inputs_check(context, 2, context.get_input_size());
+    auto x = context.get_input(0);
+    std::deque<Output<Node>> list_elems;
+    for (size_t i = 1; i < context.get_input_size(); i++) {
+        auto index = context.get_input(static_cast<int>(i));
+        if (index.get_element_type() == element::i64) {
+            auto converted = context.mark_node(std::make_shared<ov::op::v0::Convert>(index, element::i32));
+            list_elems.push_back(converted);
+        } else {
+            list_elems.push_back(index);
+        }
+    }
+    auto concat =
+        context.mark_node(std::make_shared<ov::op::v0::Concat>(OutputVector(list_elems.begin(), list_elems.end()), 0));
+    auto gather = std::make_shared<ov::op::v8::GatherND>(x, concat);
+    return {gather};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op/layer_norm.cpp

@@ -7,6 +7,7 @@
 #include "openvino/op/constant.hpp"
 #include "openvino/op/multiply.hpp"
 #include "openvino/op/mvn.hpp"
+#include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -37,7 +38,12 @@ OutputVector translate_layer_norm(const NodeContext& context) {
     return {out_node};
 };
 
+OutputVector translate_layer_norm_fx(const NodeContext& context) {
+    auto output = translate_layer_norm(context);
+    return {context.mark_node(make_list_construct(output))};
+}
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/log.cpp

@@ -8,6 +8,8 @@
 #include "openvino/op/constant.hpp"
 #include "openvino/op/convert.hpp"
 #include "openvino/op/divide.hpp"
+#include "openvino/op/exp.hpp"
+#include "openvino/op/reduce_sum.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -38,7 +40,22 @@ OutputVector translate_log2(const NodeContext& context) {
     return {res};
 };
 
+OutputVector translate_logsumexp(const NodeContext& context) {
+    num_inputs_check(context, 1, 2);
+    auto input = context.get_input(0);
+    ov::Output<ov::Node> dim;
+    if (!context.input_is_none(1)) {
+        dim = context.get_input(1);
+    } else {
+        dim = context.mark_node(get_axes_range(context, 0));
+    }
+    auto exp = context.mark_node(std::make_shared<v0::Exp>(input));
+    auto sum = context.mark_node(std::make_shared<v1::ReduceSum>(exp, dim, false));
+    auto log = context.mark_node(std::make_shared<v0::Log>(sum));
+    return {log};
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/max_poolnd.cpp

@@ -4,6 +4,7 @@
 
 #include "openvino/frontend/pytorch/node_context.hpp"
 #include "openvino/op/max_pool.hpp"
+#include "openvino/op/util/framework_node.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -14,7 +15,7 @@ namespace op {
 using namespace ov::op;
 
 OutputVector translate_max_poolnd(const NodeContext& context) {
-    num_inputs_check(context, 6, 6);
+    num_inputs_check(context, 3, 6);
     auto kernel = context.const_input<Shape>(1);
     Strides strides;
     if (!context.input_is_none(2)) {
@@ -24,15 +25,33 @@ OutputVector translate_max_poolnd(const NodeContext& context) {
         // In case strides are not provided default is kernel
         strides = kernel;
     }
-    auto pads = context.const_input<Shape>(3);  // pytorch supports only symmetric paddings
-    auto dilations = context.const_input<Strides>(4);
-    auto rounding_type = context.const_input<bool>(5) ? RoundingType::CEIL : RoundingType::FLOOR;
+    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 : RoundingType::FLOOR;
+    }
 
     return {context.mark_node(
         std::make_shared<v8::MaxPool>(context.get_input(0), strides, dilations, pads, pads, kernel, rounding_type))};
 };
 
+OutputVector translate_max_poolnd_fx(const NodeContext& context) {
+    auto output = translate_max_poolnd(context);
+    return {context.mark_node(make_list_construct(output))};
+}
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/min_max.cpp

@@ -81,4 +81,4 @@ OutputVector translate_min(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/norm.cpp

@@ -156,7 +156,7 @@ OutputVector translate_linalg_vector_norm(const NodeContext& context) {
     // dtype=None) -> Tensor
     // aten::linalg_vector_norm.out(Tensor self, Scalar ord=2, int[1]? dim=None, bool
     // keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!):
-    num_inputs_check(context, 5, 6);
+    num_inputs_check(context, 4, 6);
     auto x = context.get_input(0);
     // ord defines the vector norm that is computed.
     auto ord = context.const_input<float>(1);
@@ -321,4 +321,4 @@ OutputVector translate_frobenius_norm(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/rsub.cpp

@@ -16,18 +16,22 @@ namespace op {
 using namespace ov::op;
 
 OutputVector translate_rsub(const NodeContext& context) {
-    num_inputs_check(context, 3, 3);
+    num_inputs_check(context, 2, 3);
     auto self = context.get_input(0);
     auto other = context.get_input(1);
-    auto alpha = context.get_input(2);
+    if (!context.input_is_none(2)) {
+        auto alpha = context.get_input(2);
+        align_eltwise_input_types(context, self, other);
+        // reverse aten::sub other - self * alpha
+        auto alpha_casted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, self));
+        auto alpha_mul = context.mark_node(std::make_shared<v1::Multiply>(self, alpha_casted));
+        return {context.mark_node(std::make_shared<v1::Subtract>(other, alpha_mul))};
+    }
     align_eltwise_input_types(context, self, other);
-    // reverse aten::sub other - self * alpha
-    auto alpha_casted = context.mark_node(std::make_shared<v1::ConvertLike>(alpha, self));
-    auto alpha_mul = context.mark_node(std::make_shared<v1::Multiply>(self, alpha_casted));
-    return {context.mark_node(std::make_shared<v1::Subtract>(other, alpha_mul))};
+    return {context.mark_node(std::make_shared<v1::Subtract>(other, self))};
 };
 
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/slice.cpp

@@ -18,7 +18,7 @@ namespace op {
 
 using namespace ov::op;
 
-OutputVector translate_slice(const NodeContext& context) {
+OutputVector translate_slice_common(const NodeContext& context, const size_t num_inputs) {
     // aten::slice.t(t[] l, int? start=None, int? end=None, int step=1) -> (t[])
     // aten::slice.Tensor(Tensor(a) self, int dim=0, int? start=None, int? end=None, int step=1) -> (Tensor(a))
     ov::Output<ov::Node> dim;
@@ -26,7 +26,7 @@ OutputVector translate_slice(const NodeContext& context) {
     int end_idx;
     int step_idx;
     auto axis_0 = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0}));
-    if (context.get_input_size() == 5) {
+    if (num_inputs == 5) {
         dim = context.get_input(1);
         if (dim.get_partial_shape().rank().is_dynamic() || dim.get_partial_shape().rank().get_length() == 0) {
             dim = context.mark_node(std::make_shared<v0::Unsqueeze>(dim, axis_0));
@@ -34,7 +34,7 @@ OutputVector translate_slice(const NodeContext& context) {
         start_idx = 2;
         end_idx = 3;
         step_idx = 4;
-    } else if (context.get_input_size() == 4) {
+    } else if (num_inputs == 4) {
         start_idx = 1;
         end_idx = 2;
         step_idx = 3;
@@ -74,6 +74,16 @@ OutputVector translate_slice(const NodeContext& context) {
     return {context.mark_node(std::make_shared<v8::Slice>(context.get_input(0), start, end, step, dim))};
 };
 
+OutputVector translate_slice(const NodeContext& context) {
+    return translate_slice_common(context, context.get_input_size());
+};
+
+OutputVector translate_slice_fx(const NodeContext& context) {
+    // slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a)
+    // FX version of slice have the inputs in the same order as it has 5 inputs, even if it has less than 5 inputs
+    return translate_slice_common(context, 5);
+};
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/op/softmax.cpp

@@ -13,17 +13,27 @@ namespace pytorch {
 namespace op {
 
 using namespace ov::op;
-OutputVector translate_softmax(const NodeContext& context) {
+OutputVector translate_softmax_common(const NodeContext& context, const bool convert_dtype) {
     num_inputs_check(context, 2, 3);
     auto x = context.get_input(0);
     auto axis = context.const_input<int64_t>(1);
-    if (!context.input_is_none(2)) {
+    // Optional input 2 to translate from half(f16) to float
+    if (convert_dtype && !context.input_is_none(2)) {
         x = apply_dtype(context, 2, x);
     }
     return {context.mark_node(std::make_shared<v8::Softmax>(x, axis))};
 };
 
+OutputVector translate_softmax(const NodeContext& context) {
+    return translate_softmax_common(context, true);
+}
+
+OutputVector translate_softmax_fx(const NodeContext& context) {
+    // _softmax(Tensor self, int dim, bool half_to_float) -> Tensor
+    return translate_softmax_common(context, false);
+}
+
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op/split.cpp

@@ -0,0 +1,41 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/op/split.hpp"
+
+#include <climits>
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/util/framework_node.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_chunk_fx(const NodeContext& context) {
+    num_inputs_check(context, 2, 3);
+    auto num_chunks = context.const_input<int>(1);
+    auto dim = context.get_input(2);
+
+    std::shared_ptr<ov::Node> chunk;
+    auto dim_val = context.const_input<int>(2);
+    auto shape = context.get_input(0).get_shape();
+    if (dim_val < 0) {
+        dim_val = static_cast<int>(shape.size()) + dim_val;
+    }
+    int num_splits = static_cast<int>(shape[dim_val]) / num_chunks;
+
+    chunk = context.mark_node(std::make_shared<v1::Split>(context.get_input(0), dim, num_splits));
+
+    return {context.mark_node(make_list_construct(chunk->outputs()))};
+}
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op/transpose.cpp

@@ -63,4 +63,4 @@ OutputVector translate_t(const NodeContext& context) {
 }  // namespace op
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -15,6 +15,7 @@ namespace op {
 
 #define OP_CONVERTER(op) OutputVector op(const NodeContext& node)
 
+// TorchScript translations
 OP_CONVERTER(translate_adaptive_avg_pool3d);
 OP_CONVERTER(translate_adaptive_max_pool2d);
 OP_CONVERTER(translate_add);
@@ -22,6 +23,7 @@ OP_CONVERTER(translate_addcmul);
 OP_CONVERTER(translate_addmm);
 OP_CONVERTER(translate_all);
 OP_CONVERTER(translate_arange);
+OP_CONVERTER(translate_argmax);
 OP_CONVERTER(translate_argsort);
 OP_CONVERTER(translate_argmax);
 OP_CONVERTER(translate_argmin);
@@ -29,7 +31,9 @@ OP_CONVERTER(translate_as_tensor);
 OP_CONVERTER(translate_avg_poolnd);
 OP_CONVERTER(translate_bool);
 OP_CONVERTER(translate_batch_norm);
+OP_CONVERTER(translate_bitwise_and);
 OP_CONVERTER(translate_bitwise_not);
+OP_CONVERTER(translate_bitwise_or);
 OP_CONVERTER(translate_cat);
 OP_CONVERTER(translate_cdist);
 OP_CONVERTER(translate_clamp);
@@ -84,6 +88,7 @@ OP_CONVERTER(translate_list_construct);
 OP_CONVERTER(translate_log);
 OP_CONVERTER(translate_log_softmax);
 OP_CONVERTER(translate_log2);
+OP_CONVERTER(translate_logsumexp);
 OP_CONVERTER(translate_loop);
 OP_CONVERTER(translate_masked_fill);
 OP_CONVERTER(translate_max);
@@ -170,10 +175,24 @@ OP_CONVERTER(translate_zeros_like);
 OP_CONVERTER(translate_quantized_convnd);
 OP_CONVERTER(translate_quantized_convnd_relu);
 OP_CONVERTER(translate_quantized_linear);
+// Torch FX Translations
+OP_CONVERTER(translate_arange_fx);
+OP_CONVERTER(translate_batch_norm_fx);
+OP_CONVERTER(translate_cat_fx);
+OP_CONVERTER(translate_chunk_fx);
+OP_CONVERTER(translate_expand_fx);
+OP_CONVERTER(translate_group_norm_fx);
+OP_CONVERTER(translate_index_fx);
+OP_CONVERTER(translate_layer_norm_fx);
+OP_CONVERTER(translate_max_poolnd_fx);
+OP_CONVERTER(translate_slice_fx);
+OP_CONVERTER(translate_softmax_fx);
+OP_CONVERTER(translate_transpose_fx);
 
 }  // namespace op
 
-const std::map<std::string, CreatorFunction> get_supported_ops() {
+// Supported ops for TorchScript
+const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
     return {
         {"aten::__and__", op::translate_1to1_match_2_inputs<opset10::LogicalAnd>},  // TODO: cover numerical cases
         {"aten::__derive_index", op::translate_derive_index},
@@ -447,6 +466,104 @@ const std::map<std::string, CreatorFunction> get_supported_ops() {
     };
 };
 
+const std::map<std::string, CreatorFunction> get_supported_ops_fx() {
+    return {
+        {"aten._adaptive_avg_pool2d.default", op::translate_1to1_match_2_inputs<opset10::AdaptiveAvgPool>},
+        {"aten._softmax.default", op::translate_softmax_fx},
+        {"aten._to_copy.default", op::skip_node},
+        {"aten._unsafe_view.default", op::translate_reshape},
+        {"aten.add.Tensor", op::translate_add},
+        {"aten.add_.Tensor", op::translate_add},
+        {"aten.addmm.default", op::translate_addmm},
+        {"aten.arange.start", op::translate_arange_fx},
+        {"aten.arange.default", op::translate_arange_fx},
+        {"aten.argmax.default", op::translate_argmax},
+        {"aten.avg_pool2d.default", op::translate_avg_poolnd},
+        {"aten.bitwise_and.Tensor", op::translate_bitwise_and},
+        {"aten.bmm.default", op::translate_1to1_match_2_inputs_align_types<opset10::MatMul>},
+        {"aten.cat.default", op::translate_cat_fx},
+        {"aten.clamp_min.default", op::translate_1to1_match_2_inputs<opset10::Maximum>},
+        {"aten.clone.default", op::skip_node},  // ignore clone operators that are inserted by PyTorch autograd
+        {"aten.convolution.default", op::translate_convolution},
+        {"aten.copy_.default", op::skip_node},
+        {"aten.cos.default", op::translate_1to1_match_1_inputs<opset10::Cos>},
+        {"aten.cumsum.default", op::translate_cumsum},
+        {"aten.detach.default", op::skip_node},
+        {"aten.div.Scalar", op::translate_div},
+        {"aten.div.Tensor", op::translate_div},
+        {"aten.embedding.default", op::translate_embedding},
+        {"aten.empty.memory_format", op::translate_empty},
+        {"aten.eq.Scalar", op::translate_1to1_match_2_inputs_align_types<opset10::Equal>},
+        {"aten.eq.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Equal>},
+        {"aten.exp.default", op::translate_1to1_match_1_inputs<opset10::Exp>},
+        {"aten.expand.default", op::translate_expand_fx},
+        {"aten.full.default", op::translate_full},
+        {"aten.gather.default", op::translate_gather},
+        {"aten.gelu.default", op::translate_gelu},
+        {"aten.gt.Scalar", op::translate_1to1_match_2_inputs_align_types<opset10::Greater>},
+        {"aten.hardsigmoid.default", op::translate_1to1_match_1_inputs<opset10::HSigmoid>},
+        {"aten.hardswish_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::HSwish>>},
+        {"aten.hardtanh_.default", op::inplace_op<op::translate_hardtanh>},
+        {"aten.index.Tensor", op::translate_index_fx},
+        {"aten.lift_fresh_copy.default", op::skip_node},
+        {"aten.linalg_vector_norm.default", op::translate_linalg_vector_norm},
+        {"aten.log.default", op::translate_log},
+        {"aten.logsumexp.default", op::translate_logsumexp},
+        {"aten.lt.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Less>},
+        {"aten.masked_fill_.Scalar", op::inplace_op<op::translate_masked_fill>},
+        {"aten.masked_fill.Tensor", op::translate_masked_fill},
+        {"aten.max_pool2d_with_indices.default", op::translate_max_poolnd_fx},
+        {"aten.mean.dim", op::translate_mean},
+        {"aten.mm.default", op::translate_1to1_match_2_inputs<opset10::MatMul>},
+        {"aten.mul.Tensor", op::translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
+        {"aten.mul.Scalar", op::translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
+        {"aten.native_batch_norm.default", op::translate_batch_norm_fx},
+        {"aten.native_group_norm.default", op::translate_group_norm_fx},
+        {"aten.native_layer_norm.default", op::translate_layer_norm_fx},
+        {"aten.neg.default", op::translate_neg},
+        {"aten.new_ones.default", op::translate_new_ones},
+        {"aten.permute.default", op::translate_1to1_match_2_inputs<opset10::Transpose>},
+        {"aten.pow.Tensor_Scalar", op::translate_pow},
+        {"aten.relu.default", op::translate_1to1_match_1_inputs<opset10::Relu>},
+        {"aten.relu_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Relu>>},
+        {"aten.rsub.Scalar", op::translate_rsub},
+        {"aten.select.int", op::translate_select},
+        {"aten.sigmoid.default", op::translate_1to1_match_1_inputs<opset10::Sigmoid>},
+        {"aten.silu.default", op::translate_1to1_match_1_inputs<opset10::Swish>},
+        {"aten.silu_.default", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Swish>>},
+        {"aten.sin.default", op::translate_1to1_match_1_inputs<opset10::Sin>},
+        {"aten.slice.Tensor", op::translate_slice_fx},
+        {"aten.split.Tensor", op::translate_chunk_fx},
+        {"aten.sub.default", op::translate_sub},
+        {"aten.sub.Tensor", op::translate_sub},
+        {"aten.t.default", op::translate_t},
+        {"aten.tanh.default", op::translate_1to1_match_1_inputs<opset10::Tanh>},
+        {"aten.transpose.int", op::translate_transpose},
+        {"aten.unsqueeze.default", op::translate_1to1_match_2_inputs<opset10::Unsqueeze>},
+        {"aten.upsample_nearest2d.default", op::translate_upsample_nearest2d},
+        {"aten.view.default", op::translate_reshape},
+        {"aten.where.self", op::translate_where},
+        {"aten.zeros_like.default", op::translate_zeros_like},
+        {"get_attr", op::translate_constant},
+        {"prim::Constant", op::translate_constant},
+        {"prim::device", op::translate_constant},
+        {"prim::GetAttr", op::translate_get_attr},
+        {"prim::If", op::translate_if},
+        {"prim::is_cuda", op::return_false_scalar},
+        {"prim::ListConstruct", op::translate_list_construct},
+        {"prim::Loop", op::translate_loop},
+        {"prim::NumToTensor", op::skip_node},  // In openvino we already store number as tensor with shape []
+        {"prim::PythonOp", op::translate_pythonop},
+        {"prim::requires_grad", op::return_false_scalar},
+        {"prim::type", op::skip_node},  // Used with prim::device, pass PtFrameworkNode.
+        {"torchvision::deform_conv2d", op::translate_deform_conv},
+        {"torchvision::nms", op::translate_nms},
+        {"torchvision::roi_align", op::translate_roi_align},
+        {"<built-in function getitem>", op::translate_getitem},  // TODO: Check if there is any other way to handle this
+
+    };
+};
+
 }  // namespace pytorch
 }  // namespace frontend
 }  // namespace ov

src/frontends/pytorch/src/op_table.hpp

@@ -10,7 +10,8 @@ namespace ov {
 namespace frontend {
 namespace pytorch {
 
-const std::map<std::string, CreatorFunction> get_supported_ops();
+const std::map<std::string, CreatorFunction> get_supported_ops_ts();
+const std::map<std::string, CreatorFunction> get_supported_ops_fx();
 
 }  // namespace pytorch
 }  // namespace frontend

src/frontends/pytorch/src/transforms.cpp

@@ -380,4 +380,4 @@ void apply_pytorch_conversion_transforms(std::shared_ptr<ov::Model> model) {
 
 }  // namespace pytorch
 }  // namespace frontend
-}  // namespace ov
+}  // namespace ov

src/frontends/pytorch/src/transforms/softmax_reshape_elimination.cpp

@@ -0,0 +1,84 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "softmax_reshape_elimination.hpp"
+
+#include <ngraph/validation_util.hpp>
+
+#include "openvino/core/rt_info.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/softmax.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace pass {
+
+SoftmaxReshapeElimination::SoftmaxReshapeElimination() {
+    const auto m_reshape0 = ov::pass::pattern::wrap_type<ov::op::v1::Reshape>(ov::pass::pattern::has_static_shape());
+    const auto m_softmax = ov::pass::pattern::wrap_type<ov::op::v1::Softmax, ov::op::v8::Softmax>({m_reshape0});
+    const auto m_reshape1 = ov::pass::pattern::wrap_type<ov::op::v1::Reshape>(
+        {m_softmax, ov::pass::pattern::wrap_type<ov::op::v0::Constant>()});
+
+    register_matcher(
+        std::make_shared<ov::pass::pattern::Matcher>(m_reshape1,
+                                                     "ov::frontend::pytorch::pass::PrimTupleUnpackReplacer"),
+        [=](ov::pass::pattern::Matcher& m) {
+            auto& pattern_to_output = m.get_pattern_value_map();
+            auto reshape0 = pattern_to_output[m_reshape0].get_node_shared_ptr();
+            auto softmax = pattern_to_output[m_softmax].get_node_shared_ptr();
+            auto reshape1 = pattern_to_output[m_reshape1].get_node_shared_ptr();
+
+            const auto input_shape = reshape0->get_input_partial_shape(0);
+            const auto output_shape = reshape1->get_output_partial_shape(0);
+            if (input_shape.is_dynamic() || output_shape.is_dynamic() ||
+                input_shape.get_shape() != output_shape.get_shape())
+                return false;
+
+            const auto softmax_rank = softmax->get_input_partial_shape(0).rank();
+            int64_t axis = 0;
+            if (const auto softmax_v8 = ov::as_type_ptr<const ov::op::v8::Softmax>(softmax)) {
+                OPENVINO_SUPPRESS_DEPRECATED_START
+                axis = ov::normalize_axis(softmax->get_friendly_name(), softmax_v8->get_axis(), softmax_rank);
+                OPENVINO_SUPPRESS_DEPRECATED_END
+            } else if (const auto softmax_v1 = ov::as_type_ptr<const ov::op::v1::Softmax>(softmax)) {
+                axis = softmax_v1->get_axis();
+            } else {
+                return false;
+            }
+
+            // Supports only last axis
+            if (axis != softmax_rank.get_length() - 1)
+                return false;
+
+            // Dimensions by reduction axis should be equal
+            if (input_shape.get_shape().back() != softmax->get_input_shape(0).back())
+                return false;
+
+            // Eliminate Reshape before Softmax
+            reshape0->output(0).replace(reshape0->input_value(0));
+            copy_runtime_info(
+                {reshape0->input_value(0).get_node_shared_ptr(), reshape0->output(0).get_node_shared_ptr()},
+                reshape0->input_value(0).get_node_shared_ptr());
+
+            // Eliminate Reshape after Softmax with name saving
+            replace_output_update_name(reshape1->output(0), reshape1->input_value(0));
+
+            // update axis
+            const auto new_axis = input_shape.rank().get_length() - 1;
+            if (auto softmax_v8 = ov::as_type_ptr<ov::op::v8::Softmax>(softmax)) {
+                softmax_v8->set_axis(new_axis);
+            } else if (auto softmax_v1 = ov::as_type_ptr<ov::op::v1::Softmax>(softmax)) {
+                softmax_v1->set_axis(new_axis);
+            }
+
+            return true;
+        });
+}
+
+}  // namespace pass
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/transforms/softmax_reshape_elimination.hpp

@@ -0,0 +1,28 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/graph_rewrite.hpp"
+#include "openvino/pass/pattern/matcher.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace pass {
+
+/**
+ * @interface SoftmaxReshapeElimination
+ * @brief The pass removes Reshape operations around Softmax if possible
+ * @ingroup snippets
+ */
+class SoftmaxReshapeElimination : public ov::pass::MatcherPass {
+public:
+    SoftmaxReshapeElimination();
+};
+
+}  // namespace pass
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/translate_session.cpp

@@ -80,6 +80,10 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
     std::shared_ptr<TorchDecoder> pytorch_model,
     const TensorMap& external_tensor_map,
     const std::unordered_map<size_t, PlaceDesc>& external_descriptors) {
+    // FIXME: don't use global variable to count inlined inputs, no we should use global counter because ID should be
+    // unique for all new inlined inputs
+    static size_t inlined_nodes_counter = 100000000;  // Suppose there are not graph with more than 10M nodes
+
     std::shared_ptr<Model> resulting_model;  // define here to make a conversion in a nested scope
     {
         auto parameters = std::make_shared<ParameterVector>();
@@ -137,6 +141,18 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
             // But this value can be found in the outer scope, for this purpose we create new input for the model to
             // link with external scope on a higher level.
 
+            auto inlined_inputs = node->inlined_inputs(inlined_nodes_counter);
+            for (size_t i = 0; i < inlined_inputs.size(); ++i) {
+                size_t fw_tensor_id = inlined_nodes_counter;
+                ++inlined_nodes_counter;  // TODO: Make sure that Decoder side use the same increment for producing ids
+                if (tensor_map->find(fw_tensor_id) != tensor_map->end()) {
+                    throw std::runtime_error("Duplicated producer for PT value with unique ID: " +
+                                             std::to_string(fw_tensor_id) + " produced by inlined_inputs");
+                }
+
+                (*tensor_map)[fw_tensor_id] = inlined_inputs[i];
+            }
+
             auto raw_inputs = node->inputs();
             for (size_t i = 0; i < raw_inputs.size(); ++i) {
                 auto input = raw_inputs.at(i);
@@ -147,7 +163,13 @@ std::shared_ptr<Model> TranslateSession::convert_pytorch_model(
                     PartialShape ps = node->get_input_shape(i);
                     auto type = simplified_type_interpret(node->get_input_type(i));
                     // TODO: Use special API to set custom type specification
-                    auto parameter = std::make_shared<v0::Parameter>(element::dynamic, ps);
+                    std::shared_ptr<v0::Parameter> parameter;
+                    // TODO: Use decoder type or explore adding the missing cast types to Torchscript path
+                    const char* torch_tracing_mode = std::getenv("PYTORCH_TRACING_MODE");
+                    if ((torch_tracing_mode != nullptr) && std::strcmp(torch_tracing_mode, "TORCHFX") == 0)
+                        parameter = std::make_shared<v0::Parameter>(type.as<element::Type>(), ps);
+                    else
+                        parameter = std::make_shared<v0::Parameter>(element::dynamic, ps);
                     // TODO: Missing get_input_transpose_order handling for not trivial layouts
                     (*tensor_map)[input] = parameter;
                     // set name of parameter to the index of node in the model

src/frontends/pytorch/src/utils.cpp

@@ -336,6 +336,16 @@ std::shared_ptr<ov::op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node>
     return fw_node;
 }
 
+std::shared_ptr<ov::op::util::FrameworkNode> make_list_construct(const ov::OutputVector& inputs) {
+    auto list_construct = std::make_shared<::ov::op::util::FrameworkNode>(inputs, inputs.size());
+    ov::op::util::FrameworkNodeAttrs attrs;
+    attrs.set_type_name("PTFrameworkNode");
+    attrs[PtFrameworkNode::op_type_key] = "prim::ListConstruct";
+    list_construct->set_attrs(attrs);
+    list_construct->validate_and_infer_types();
+    return list_construct;
+}
+
 bool is_none_node(const Output<Node>& node) {
     if (const auto& fw_node_inp = std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(node.get_node_shared_ptr())) {
         const auto& attrs = fw_node_inp->get_attrs();

src/frontends/pytorch/src/utils.hpp

@@ -58,6 +58,8 @@ OutputVector make_framework_node(const NodeContext& context, const std::string&
 
 std::shared_ptr<op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node, const std::string& type);
 
+std::shared_ptr<op::util::FrameworkNode> make_list_construct(const ov::OutputVector& inputs);
+
 bool is_none_node(const Output<Node>& node);
 
 // TODO: Eliminate the need of this function by implementing more accurate custom data type handling

tests/layer_tests/py_frontend_tests/test_torch_decoder.py

@@ -37,7 +37,7 @@ def get_traced_model(model, inputs=[], frozen=True):
 
 @pytest.mark.precommit
 def test_pytorch_decoder_get_output_type_str():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.frontend.pytorch.py_pytorch_frontend import _Type as DecoderType
 
     model = get_scripted_model(AtenDiv("trunc"))
@@ -53,7 +53,7 @@ def test_pytorch_decoder_get_output_type_str():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_get_output_type_none():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.frontend.pytorch.py_pytorch_frontend import _Type as DecoderType
 
     model = get_scripted_model(AtenDiv(None))
@@ -69,7 +69,7 @@ def test_pytorch_decoder_get_output_type_none():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_get_input_type_str():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.frontend.pytorch.py_pytorch_frontend import _Type as DecoderType
 
     model = get_scripted_model(AtenDiv("trunc"))
@@ -83,7 +83,7 @@ def test_pytorch_decoder_get_input_type_str():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_get_input_type_none():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.frontend.pytorch.py_pytorch_frontend import _Type as DecoderType
 
     model = get_scripted_model(AtenDiv(None))
@@ -97,7 +97,7 @@ def test_pytorch_decoder_get_input_type_none():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp16_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -119,7 +119,7 @@ def test_pytorch_decoder_can_convert_fp16_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bf16_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -141,7 +141,7 @@ def test_pytorch_decoder_can_convert_bf16_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp32_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -163,7 +163,7 @@ def test_pytorch_decoder_can_convert_fp32_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_fp64_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -185,7 +185,7 @@ def test_pytorch_decoder_can_convert_fp64_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -207,7 +207,7 @@ def test_pytorch_decoder_can_convert_bool_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_u8_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -229,7 +229,7 @@ def test_pytorch_decoder_can_convert_u8_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i8_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -251,7 +251,7 @@ def test_pytorch_decoder_can_convert_i8_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i32_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -273,7 +273,7 @@ def test_pytorch_decoder_can_convert_i32_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_i64_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -295,7 +295,7 @@ def test_pytorch_decoder_can_convert_i64_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int64_max():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     class I64MaxConst(torch.nn.Module):
         def forward(self):
@@ -313,7 +313,7 @@ def test_pytorch_decoder_can_convert_int64_max():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_list():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -336,7 +336,7 @@ def test_pytorch_decoder_can_convert_int_list():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_list():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -359,7 +359,7 @@ def test_pytorch_decoder_can_convert_float_list():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_list():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -382,7 +382,7 @@ def test_pytorch_decoder_can_convert_bool_list():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_tuple():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -405,7 +405,7 @@ def test_pytorch_decoder_can_convert_int_tuple():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_tuple():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -428,7 +428,7 @@ def test_pytorch_decoder_can_convert_float_tuple():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_bool_tuple():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class ListConst(torch.nn.Module):
@@ -451,7 +451,7 @@ def test_pytorch_decoder_can_convert_bool_tuple():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_empty_list():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class aten_roll(torch.nn.Module):
@@ -479,7 +479,7 @@ def test_pytorch_decoder_can_convert_empty_list():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_int_scalar_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -510,7 +510,7 @@ def test_pytorch_decoder_can_convert_int_scalar_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_float_scalar_tensor():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
 
     class SomeTensor(torch.nn.Module):
@@ -541,7 +541,7 @@ def test_pytorch_decoder_can_convert_float_scalar_tensor():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_tensor_list():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from openvino.runtime import PartialShape, Type
     from typing import List, Optional
 
@@ -579,7 +579,7 @@ def test_pytorch_decoder_can_convert_tensor_list():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_tensor_list_empty():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from typing import List, Optional
 
     class SomeTensor(torch.nn.Module):
@@ -607,7 +607,7 @@ def test_pytorch_decoder_can_convert_tensor_list_empty():
 
 @pytest.mark.precommit
 def test_pytorch_decoder_can_convert_optional_tensor_none():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     from typing import Optional
 
     class SomeTensor(torch.nn.Module):

tests/layer_tests/py_frontend_tests/test_torch_frontend.py

@@ -30,7 +30,7 @@ def get_scripted_model(model):
 
 
 def test_pytorch_fe_set_input_shape():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     model = get_scripted_model(aten_relu())
     decoder = TorchScriptPythonDecoder(model)
@@ -45,7 +45,7 @@ def test_pytorch_fe_set_input_shape():
 
 
 def test_pytorch_fe_set_input_type():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     model = get_scripted_model(aten_relu())
     decoder = TorchScriptPythonDecoder(model)
@@ -59,7 +59,7 @@ def test_pytorch_fe_set_input_type():
 
 
 def test_pytorch_fe_set_input_value():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     model = get_scripted_model(aten_relu())
     decoder = TorchScriptPythonDecoder(model)
@@ -75,7 +75,7 @@ def test_pytorch_fe_set_input_value():
 
 
 def test_conversion_extension():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     class Model(torch.nn.Module):
         def __init__(self):
@@ -185,7 +185,7 @@ def get_builtin_extensions_path():
 
 
 def test_op_extension():
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     class Elu(torch.nn.Module):
         def __init__(self, alpha):
@@ -216,7 +216,7 @@ def test_op_extension():
 
 def test_pytorch_telemetry():
     from openvino.frontend import TelemetryExtension
-    from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+    from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
     class MockTelemetry:
         def __init__(self, stat):

tests/layer_tests/pytest.ini

@@ -3,4 +3,5 @@ markers =
     nightly
     precommit
     precommit_ts_backend
+    precommit_fx_backend
     timeout

tests/layer_tests/pytorch_tests/pytorch_layer_test_class.py

@@ -8,7 +8,7 @@ import os
 
 import numpy as np
 from common.constants import test_device, test_precision
-from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
 from openvino.frontend import FrontEndManager
 from openvino.runtime import Core, Type, PartialShape
@@ -65,13 +65,17 @@ class PytorchLayerTest:
         else:
             custom_eps = 1e-4
 
-        def use_ts_backend():
-            return(os.environ.get('USE_TS_BACKEND', False))
+        def use_torch_compile_backend():
+            torch_compile_env = os.getenv("PYTORCH_TRACING_MODE")
+            if torch_compile_env is not None:
+                if (torch_compile_env == "TORCHFX" or torch_compile_env == "TORCHSCRIPT"):
+                    return True
+            return False
 
         ov_inputs = flattenize_inputs(inputs)
 
-        if use_ts_backend():
-            self.ts_backend_test(model, torch_inputs, custom_eps)
+        if use_torch_compile_backend():
+            self.torch_compile_backend_test(model, torch_inputs, custom_eps)
         else:
             with torch.no_grad():
                 model.eval()
@@ -218,7 +222,7 @@ class PytorchLayerTest:
         om.validate_nodes_and_infer_types()
         return om
 
-    def ts_backend_test(self, model, inputs, custom_eps):
+    def torch_compile_backend_test(self, model, inputs, custom_eps):
         torch._dynamo.reset()
         with torch.no_grad():
             model.eval()

tests/layer_tests/pytorch_tests/test_adaptive_avg_pool3d.py

@@ -33,6 +33,7 @@ class TestAdaptiveAvgPool3D(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_adaptive_avg_pool3d(self, ie_device, precision, ir_version, input_tensor, output_size):
         self.input_tensor = input_tensor
         self._test(*self.create_model(output_size), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_adaptive_max_pool_2d.py

@@ -47,6 +47,7 @@ class TestAdaptiveMaxPool2D(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_adaptive_max_pool2d(self, ie_device, precision, ir_version, input_tensor, output_size, return_indices):
         self.input_tensor = input_tensor
         self._test(*self.create_model(output_size, return_indices), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_add.py

@@ -39,6 +39,7 @@ class TestAdd(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     @pytest.mark.parametrize("op_type", ["add", "add_"])
     def test_add(self, ie_device, precision, ir_version, alpha, input_rhs, op_type):
         self.input_rhs = input_rhs
@@ -100,6 +101,7 @@ class TestAddTypes(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_add_types(self, ie_device, precision, ir_version, lhs_type, lhs_shape, rhs_type, rhs_shape):
         self.lhs_type = lhs_type
         self.lhs_shape = lhs_shape
@@ -123,5 +125,6 @@ class TestAddLists(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_add(self, ie_device, precision, ir_version):
         self._test(*self.create_model(), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_batch_norm.py

@@ -58,6 +58,7 @@ class TestBatchNorm(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_batch_norm(self, weights, bias, eps, train, running_stats, ie_device, precision, ir_version, kwargs_to_prepare_input):
         self._test(*self.create_model(weights, bias, eps, train, running_stats),
                    ie_device, precision, ir_version, kwargs_to_prepare_input=kwargs_to_prepare_input, dynamic_shapes=False, use_mo_convert=False)

tests/layer_tests/pytorch_tests/test_convnd.py

@@ -3,7 +3,7 @@
 
 import pytest
 from openvino.frontend import FrontEndManager
-from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 from pytorch_layer_test_class import PytorchLayerTest
 
 
@@ -165,7 +165,7 @@ class TestConv2DInSubgraph(PytorchLayerTest):
     def _prepare_input(self):
         import numpy as np
         return (np.random.randn(2, 3, 25, 25).astype(np.float32), np.array([1], dtype=np.int32))
-    
+
     def convert_directly_via_frontend(self, model, example_input, trace_model, dynamic_shapes, ov_inputs, freeze_model):
         # Overload function to allow reproduction of issue caused by additional freeze.
         import torch

tests/layer_tests/pytorch_tests/test_convolution.py

@@ -208,6 +208,7 @@ class TestConvolution(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_convolution1d(self, params, bias, underscore, ie_device, precision, ir_version):
         self._test(*self.create_model(**params, bias=bias, underscore=underscore),
                    ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1,
@@ -219,6 +220,7 @@ class TestConvolution(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_convolution2d(self, params, bias, underscore, ie_device, precision, ir_version):
         self._test(*self.create_model(**params, bias=bias, underscore=underscore),
                    ie_device, precision, ir_version, dynamic_shapes=params['groups'] == 1)

tests/layer_tests/pytorch_tests/test_quantized_convnd.py

@@ -6,7 +6,7 @@ import numpy as np
 import torch
 
 from openvino.frontend import FrontEndManager
-from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 from pytorch_layer_test_class import PytorchLayerTest
 
 

tests/layer_tests/pytorch_tests/test_remainder.py

@@ -97,6 +97,7 @@ class TestRemainderTypes(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     @pytest.mark.precommit_ts_backend
+    @pytest.mark.precommit_fx_backend
     def test_remainder_types(self, ie_device, precision, ir_version, lhs_type, lhs_shape, rhs_type, rhs_shape):
         self.lhs_type = lhs_type
         self.lhs_shape = lhs_shape

tools/mo/openvino/tools/mo/convert_impl.py

@@ -530,8 +530,8 @@ def check_model_object(argv):
         if isinstance(model, (torch.nn.Module, torch.jit.ScriptFunction)):
             return "pytorch"
         try:
-            from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
-            
+            from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
+
             if isinstance(model, TorchScriptPythonDecoder):
                 return "pytorch"
         except Exception as e:
@@ -802,7 +802,7 @@ def update_args_for_saved_model_dir(args: dict):
             'input_model' in args and args['input_model'] is not None:
         raise Error("Both --input_model and --saved_model_dir are defined. "
                     "Please specify either input_model or saved_model_dir directory.")
-    
+
     if 'input_model' in args and isinstance(args['input_model'], (str, Path)) and os.path.isdir(args['input_model']):
         args['saved_model_dir'] = args['input_model']
         args['input_model'] = None

tools/mo/openvino/tools/mo/moc_frontend/pytorch_frontend_utils.py

@@ -15,7 +15,7 @@ from openvino.tools.mo.moc_frontend.shape_utils import get_static_shape
 
 def get_pytorch_decoder(model, input_shape, example_inputs, args):
     try:
-        from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+        from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     except Exception as e:
         log.error("PyTorch frontend loading failed")
         raise e
@@ -57,7 +57,7 @@ def get_value_from_list_or_dict(container, name, idx):
 
 def extract_input_info_from_example(args, inputs):
     try:
-        from openvino.frontend.pytorch.decoder import pt_to_ov_type_map # pylint: disable=no-name-in-module,import-error
+        from openvino.frontend.pytorch.utils import pt_to_ov_type_map  # pylint: disable=no-name-in-module,import-error
     except Exception as e:
         log.error("PyTorch frontend loading failed")
         raise e

tools/ovc/openvino/tools/ovc/convert_impl.py

@@ -162,7 +162,7 @@ def check_model_object(argv):
         if isinstance(model, (torch.nn.Module, torch.jit.ScriptFunction)):
             return "pytorch"
         try:
-            from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+            from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
 
             if isinstance(model, TorchScriptPythonDecoder):
                 return "pytorch"

tools/ovc/openvino/tools/ovc/moc_frontend/pytorch_frontend_utils.py

@@ -15,7 +15,7 @@ from openvino.tools.ovc.moc_frontend.shape_utils import get_static_shape
 
 def get_pytorch_decoder(model, example_inputs, args):
     try:
-        from openvino.frontend.pytorch.decoder import TorchScriptPythonDecoder
+        from openvino.frontend.pytorch.ts_decoder import TorchScriptPythonDecoder
     except Exception as e:
         log.error("PyTorch frontend loading failed")
         raise e
@@ -56,7 +56,7 @@ def get_value_from_list_or_dict(container, name, idx):
 
 def extract_input_info_from_example(args, inputs):
     try:
-        from openvino.frontend.pytorch.decoder import pt_to_ov_type_map # pylint: disable=no-name-in-module,import-error
+        from openvino.frontend.pytorch.utils import pt_to_ov_type_map # pylint: disable=no-name-in-module,import-error
     except Exception as e:
         log.error("PyTorch frontend loading failed")
         raise e