[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>
This commit is contained in:
Luca Tam 2024-03-08 10:55:47 +01:00 committed by GitHub
parent b068ea2a33
commit c83c06cf27
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 104 additions and 0 deletions

View File

@ -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

View File

@ -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

View File

@ -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"}
)