[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:
parent
b068ea2a33
commit
c83c06cf27
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
)
|
||||
Loading…
Reference in New Issue