From c83c06cf279fec894da26fbf6348da423d464f83 Mon Sep 17 00:00:00 2001
From: Luca Tam <155265433+LucaTamSapienza@users.noreply.github.com>
Date: Fri, 8 Mar 2024 10:55:47 +0100
Subject: [PATCH] [PT FE] aten::dot operation (#22978)

### Details:
- added `aten::dot` operation for pytorch models

### Tickets:
 - Closes #22074

---------

Co-authored-by: Maxim Vafin <maxim.vafin@intel.com>
---
 src/frontends/pytorch/src/op/dot.cpp        | 37 ++++++++++++
 src/frontends/pytorch/src/op_table.cpp      |  2 +
 tests/layer_tests/pytorch_tests/test_dot.py | 65 +++++++++++++++++++++
 3 files changed, 104 insertions(+)
 create mode 100644 src/frontends/pytorch/src/op/dot.cpp
 create mode 100644 tests/layer_tests/pytorch_tests/test_dot.py

diff --git a/src/frontends/pytorch/src/op/dot.cpp b/src/frontends/pytorch/src/op/dot.cpp
new file mode 100644
index 00000000000..2c7bf9e4ac1
--- /dev/null
+++ b/src/frontends/pytorch/src/op/dot.cpp
@@ -0,0 +1,37 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/matmul.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_dot(const NodeContext& context) {
+    // "aten::dot(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)"
+
+    auto tensor1 = context.get_input(0);
+    auto tensor2 = context.get_input(1);
+    align_eltwise_input_types(context, tensor1, tensor2, true);
+
+    OutputVector dot_product;
+    if (!context.input_is_none(2)) {
+        dot_product = {context.mark_node(std::make_shared<v0::MatMul>(tensor1, tensor2))};
+        context.mutate_input(2, dot_product[0]);
+    } else {
+        dot_product = translate_1to1_match_2_inputs<v0::MatMul>(context);
+    }
+
+    return dot_product;
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov
diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp
index d846badb5ca..ae62d4b30e7 100644
--- a/src/frontends/pytorch/src/op_table.cpp
+++ b/src/frontends/pytorch/src/op_table.cpp
@@ -64,6 +64,7 @@ OP_CONVERTER(translate_derive_index);
 OP_CONVERTER(translate_dim);
 OP_CONVERTER(translate_div);
 OP_CONVERTER(translate_div_);
+OP_CONVERTER(translate_dot);
 OP_CONVERTER(translate_elu);
 OP_CONVERTER(translate_embedding);
 OP_CONVERTER(translate_embedding_bag);
@@ -393,6 +394,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::dim", op::translate_dim},
         {"aten::div", op::translate_div},
         {"aten::div_", op::translate_div_},
+        {"aten::dot", op::translate_dot},
         {"aten::dropout", op::skip_node},
         {"aten::dropout_", op::skip_node},
         // aten::einsum - Supported in limited set of patterns
diff --git a/tests/layer_tests/pytorch_tests/test_dot.py b/tests/layer_tests/pytorch_tests/test_dot.py
new file mode 100644
index 00000000000..290e1f052d5
--- /dev/null
+++ b/tests/layer_tests/pytorch_tests/test_dot.py
@@ -0,0 +1,65 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+class TestDot(PytorchLayerTest):
+    def _prepare_input(self, inputs, dtype, out=False):
+        import numpy as np
+        x = np.array(inputs[0]).astype(dtype)
+        y = np.array(inputs[1]).astype(dtype)
+        if not out:
+            return (x, y)
+        return (x, y, np.array(0).astype(dtype))
+
+    def create_model(self, mode, dtype):
+        import torch
+
+        dtype_map = {
+            "float32": torch.float32,
+            "float64": torch.float64,
+            "int64": torch.int64,
+            "int32": torch.int32,
+            "uint8": torch.uint8,
+            "int8": torch.int8,
+        }
+
+        class aten_dot(torch.nn.Module):
+            def __init__(self, mode, dtype):
+                super().__init__()
+                self.dtype = dtype
+                if mode =="out":
+                    self.forward = self.forward_out
+                else:
+                    self.forward = self.forward_default
+
+            def forward_default(self, tensor1, tensor2):
+                return torch.dot(tensor1.to(self.dtype), tensor2.to(self.dtype))
+
+            def forward_out(self, tensor1, tensor2, y):
+                return torch.dot(tensor1.to(self.dtype), tensor2.to(self.dtype), out=y), y
+
+        dtype = dtype_map.get(dtype)
+
+        ref_net = None
+
+        return aten_dot(mode, dtype), ref_net, "aten::dot"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("mode, dtype", [
+        ("", "float32"), ("", "float64"), ("", "int32"), ("", "int64"), ("", "int8"),
+        ("out", "float32"), ("out", "float64"), ("out", "int32"), ("out", "int64"), ("out", "int8")])
+    @pytest.mark.parametrize(
+        "inputs", [([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]), ([1, 2, 3], [4, 5, 6]), ([1, 1, 1], [1, 1, 1])]
+    )
+    def test_dot(self, mode, dtype, inputs, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(mode, dtype),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={"inputs": inputs, "dtype": dtype, "out": mode == "out"}
+        )