[PT FE]: enable dtype in softmax and constant filling, extend logical ops support (#16276)

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

@@ -0,0 +1,29 @@
+// 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(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/cumsum.cpp

@@ -3,11 +3,7 @@
 //
 
 #include "openvino/frontend/pytorch/node_context.hpp"
-#include "openvino/op/constant.hpp"
-#include "openvino/op/convert.hpp"
-#include "openvino/op/convert_like.hpp"
 #include "openvino/op/cum_sum.hpp"
-#include "pt_framework_node.hpp"
 #include "utils.hpp"
 
 namespace ov {
@@ -23,15 +19,7 @@ OutputVector translate_cumsum(NodeContext& context) {
     auto x = context.get_input(0);
     auto dim = context.get_input(1);
     if (!context.input_is_none(2)) {
-        if (std::dynamic_pointer_cast<v0::Constant>(context.get_input_from_visible_context(2).get_node_shared_ptr())) {
-            auto dtype = convert_dtype(context.const_input<int64_t>(2));
-            x = context.mark_node(std::make_shared<v0::Convert>(x, dtype));
-        } else if (const auto& fw_node = cast_fw_node(context.get_input(2).get_node_shared_ptr(), "prim::dtype")) {
-            auto out_tensor = fw_node->input_value(0);
-            x = context.mark_node(std::make_shared<v1::ConvertLike>(x, out_tensor));
-        } else {
-            FRONT_END_OP_CONVERSION_CHECK(false, "Couldn't get dtype input");
-        }
+        x = apply_dtype(context, 2, x);
     }
     auto result = context.mark_node(std::make_shared<v0::CumSum>(x, dim));
     if (!context.input_is_none(3)) {

src/frontends/pytorch/src/op/full.cpp

@@ -5,7 +5,6 @@
 #include "openvino/frontend/pytorch/node_context.hpp"
 #include "openvino/op/broadcast.hpp"
 #include "openvino/op/constant.hpp"
-#include "openvino/op/convert.hpp"
 #include "openvino/op/convert_like.hpp"
 #include "openvino/op/shape_of.hpp"
 #include "utils.hpp"
@@ -22,18 +21,6 @@ Output<Node> base_translate_full(const NodeContext& context, const Output<Node>&
     return context.mark_node(std::make_shared<v3::Broadcast>(value, sizes));
 }
 
-Output<Node> base_translate_full_with_convert(const NodeContext& context,
-                                              const Output<Node>& sizes,
-                                              const Output<Node>& value,
-                                              size_t dtype_id) {
-    auto filled_tensor = base_translate_full(context, sizes, value);
-    if (!context.input_is_none(dtype_id)) {
-        auto dtype = convert_dtype(context.const_input<int64_t>(dtype_id));
-        filled_tensor = context.mark_node(std::make_shared<v0::Convert>(filled_tensor, dtype));
-    }
-    return filled_tensor;
-}
-
 Output<Node> base_translate_full_with_convertlike(const NodeContext& context,
                                                   const Output<Node>& sizes,
                                                   const Output<Node>& value,
@@ -41,6 +28,18 @@ Output<Node> base_translate_full_with_convertlike(const NodeContext& context,
     auto filled_tensor = base_translate_full(context, sizes, value);
     return context.mark_node(std::make_shared<v1::ConvertLike>(filled_tensor, out));
 }
+
+Output<Node> base_translate_full_with_convert(const NodeContext& context,
+                                              const Output<Node>& sizes,
+                                              Output<Node> value,
+                                              size_t dtype_id) {
+    if (!context.input_is_none(dtype_id)) {
+        value = apply_dtype(context, dtype_id, value);
+    }
+
+    auto filled_tensor = base_translate_full(context, sizes, value);
+    return filled_tensor;
+}
 }  // namespace
 
 OutputVector translate_full(NodeContext& context) {
@@ -172,7 +171,10 @@ OutputVector translate_new_ones(NodeContext& context) {
 };
 
 OutputVector translate_empty(NodeContext& context) {
-    num_inputs_check(context, 1, 5);
+    // aten::empty(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool?
+    // pin_memory=None, MemoryFormat? memory_format=None) -> Tensor layout, device and work with memory ignored on our
+    // side, so just skip these parameters
+    num_inputs_check(context, 1, 6);
     auto sizes = context.get_input(0);
     // In OV uninitialised data is not supported, so we create a tensor filled with zeros with a given shape and type.
     auto value = context.mark_node(v0::Constant::create(element::f32, Shape{}, {0}));

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

@@ -12,11 +12,15 @@ namespace frontend {
 namespace pytorch {
 namespace op {
 
+using namespace ov::op;
 OutputVector translate_softmax(NodeContext& context) {
-    num_inputs_check(context, 2, 2);
+    num_inputs_check(context, 2, 3);
     auto x = context.get_input(0);
     auto axis = context.const_input<int64_t>(1);
-    return {context.mark_node(std::make_shared<ov::op::v8::Softmax>(x, axis))};
+    if (!context.input_is_none(2)) {
+        x = apply_dtype(context, 2, x);
+    }
+    return {context.mark_node(std::make_shared<v8::Softmax>(x, axis))};
 };
 
 }  // namespace op

src/frontends/pytorch/src/op_table.cpp

@@ -25,6 +25,7 @@ OP_CONVERTER(translate_as_tensor);
 OP_CONVERTER(translate_avg_poolnd);
 OP_CONVERTER(translate_bool);
 OP_CONVERTER(translate_batch_norm);
+OP_CONVERTER(translate_bitwise_not);
 OP_CONVERTER(translate_cat);
 OP_CONVERTER(translate_clamp);
 OP_CONVERTER(translate_constant);
@@ -134,6 +135,7 @@ const std::map<std::string, PytorchCreatorFunction> get_supported_ops() {
         {"aten::__and__", op::translate_1to1_match_2_inputs<opset10::LogicalAnd>},  // TODO: cover numerical cases
         {"aten::__getitem__", op::translate_getitem},
         {"aten::__not__", op::translate_1to1_match_1_inputs<opset10::LogicalNot>},
+        {"aten::__or__", op::translate_1to1_match_2_inputs<opset10::LogicalOr>},
         {"aten::_convolution", op::translate_convolution},
         {"aten::_convolution_mode", op::translate_convolution_mode},
         {"aten::_set_item", op::translate_set_item},
@@ -163,7 +165,9 @@ const std::map<std::string, PytorchCreatorFunction> get_supported_ops() {
         {"aten::avg_pool1d", op::translate_avg_poolnd},
         {"aten::avg_pool2d", op::translate_avg_poolnd},
         {"aten::avg_pool3d", op::translate_avg_poolnd},
+        {"aten::baddbmm", op::translate_addmm},
         {"aten::batch_norm", op::translate_batch_norm},
+        {"aten::bitwise_not", op::translate_bitwise_not},
         {"aten::bmm", op::translate_1to1_match_2_inputs<opset10::MatMul>},
         {"aten::Bool", op::translate_bool},
         {"aten::cat", op::translate_cat},

src/frontends/pytorch/src/utils.cpp

@@ -142,6 +142,21 @@ element::Type convert_dtype(int64_t pt_type) {
     return TORCH_TO_OV_TYPE.at(pt_type);
 };
 
+Output<Node> apply_dtype(const NodeContext& context, size_t dtype_port, const Output<Node>& input_tensor) {
+    if (std::dynamic_pointer_cast<opset10::Constant>(
+            context.get_input_from_visible_context(dtype_port).get_node_shared_ptr())) {
+        auto dtype = convert_dtype(context.const_input<int64_t>(dtype_port));
+        return context.mark_node(std::make_shared<opset10::Convert>(input_tensor, dtype));
+    } else if (const auto& fw_node =
+                   cast_fw_node(context.get_input(static_cast<int>(dtype_port)).get_node_shared_ptr(), "prim::dtype")) {
+        auto out_tensor = fw_node->input_value(0);
+        return context.mark_node(std::make_shared<opset10::ConvertLike>(input_tensor, out_tensor));
+    } else {
+        FRONT_END_OP_CONVERSION_CHECK(false, "Couldn't get dtype input");
+    }
+    return input_tensor;
+};
+
 ov::op::PadType convert_pad(const std::string& pt_pad) {
     FRONT_END_OP_CONVERSION_CHECK(TORCH_AUTO_PAD_TO_OV.count(pt_pad), "Unknown pad: ", pt_pad);
     return TORCH_AUTO_PAD_TO_OV.at(pt_pad);

src/frontends/pytorch/src/utils.hpp

@@ -41,6 +41,9 @@ std::shared_ptr<Node> get_axes_range(const NodeContext& context, int input_id);
 std::shared_ptr<Node> numel(const NodeContext& context, const Output<Node>& x);
 
 element::Type convert_dtype(int64_t dtype_value);
+
+Output<Node> apply_dtype(const NodeContext& context, size_t dtype_port, const Output<Node>& input_tensor);
+
 op::PadType convert_pad(const std::string& pt_pad);
 
 std::shared_ptr<Node> concat_list_construct(std::shared_ptr<Node> input);

tests/layer_tests/pytorch_tests/test_addmm.py

@@ -47,3 +47,46 @@ class TestAddMM(PytorchLayerTest):
     def test_addmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version):
         self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version,
                    kwargs_to_prepare_input=kwargs_to_prepare_input)
+
+
+class TestBAddBMM(PytorchLayerTest):
+    def _prepare_input(self, input_shape=(2, 2), matrix1_shape=(2, 2), matrix2_shape=(2, 2)):
+        import numpy as np
+        return (
+            np.random.randn(*input_shape).astype(np.float32),
+            np.random.randn(*matrix1_shape).astype(np.float32),
+            np.random.randn(*matrix2_shape).astype(np.float32)
+        )
+
+    def create_model(self, alpha, beta):
+        import torch
+
+        class aten_addmm(torch.nn.Module):
+            def __init__(self, alpha, beta):
+                super(aten_addmm, self).__init__()
+                self.alpha = alpha
+                self.beta = beta
+
+            def forward(self, m0, m1, m2):
+                return torch.baddbmm(m0, m1, m2, alpha=self.alpha, beta=self.beta)
+
+        ref_net = None
+
+        return aten_addmm(alpha, beta), ref_net, 'aten::baddbmm'
+
+    @pytest.mark.parametrize("kwargs_to_prepare_input", [
+        {"input_shape": (2, 3, 3), 'matrix1_shape': (2, 3, 3), 'matrix2_shape': (2, 3, 3)},
+        {"input_shape": (2, 2, 2), 'matrix1_shape': (2, 2, 3), 'matrix2_shape': (2, 3, 2)},
+        {"input_shape": (1, 10, 1), 'matrix1_shape': (1, 10, 5), 'matrix2_shape': (1, 5, 1)},
+        {"input_shape": (5, 1, 2), 'matrix1_shape': (5, 1, 10), 'matrix2_shape': (5, 10, 2)},
+        {"input_shape": (1, 1, 1), 'matrix1_shape': (1, 1, 10), 'matrix2_shape': (1, 10, 1)},
+
+    ])
+    @pytest.mark.parametrize("alpha,beta",
+                             [(1., 1.), (0., 1.), (1., 0.), (1., 2.), (2., 1.), (-5., -6.), (3., 4.), (0.5, 0.75),
+                              (1, 1)])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_baddbmm(self, kwargs_to_prepare_input, alpha, beta, ie_device, precision, ir_version):
+        self._test(*self.create_model(alpha, beta), ie_device, precision, ir_version,
+                   kwargs_to_prepare_input=kwargs_to_prepare_input)

tests/layer_tests/pytorch_tests/test_bitwise_not.py

@@ -0,0 +1,29 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestBitwiseNot(PytorchLayerTest):
+    def _prepare_input(self):
+        import numpy as np
+        return ((np.random.randn(1, 5) > 0).astype(bool),)
+
+    def create_model(self):
+        import torch
+
+        class aten_bitwise_not(torch.nn.Module):
+
+            def forward(self, x):
+                return torch.bitwise_not(x)
+
+        ref_net = None
+
+        return aten_bitwise_not(), ref_net, "aten::bitwise_not"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_bitwise_not(self, ie_device, precision, ir_version):
+        self._test(*self.create_model(), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_softmax.py

@@ -7,28 +7,59 @@ from pytorch_layer_test_class import PytorchLayerTest
 
 
 class TestSoftmax(PytorchLayerTest):
-    def _prepare_input(self):
+    def _prepare_input(self, second_input_dtype=None):
         import numpy as np
-        return (np.random.randn(1, 3, 224, 224).astype(np.float32),)
+        if second_input_dtype is None:
+            return (np.random.randn(1, 3, 224, 224).astype(np.float32),)
+        return (np.random.randn(1, 3, 224, 224).astype(np.float32), np.random.randn(1).astype(second_input_dtype))
 
-    def create_model(self, dim):
+    def create_model(self, dim, dtype=None, use_prim_dtype=False):
         import torch
         import torch.nn.functional as F
+        dtype_map = {
+            "float32": torch.float32,
+            "float64": torch.float64,
+            "int32": torch.int32
+        }
+        dtype = dtype_map.get(dtype)
 
         class aten_softmax(torch.nn.Module):
-            def __init__(self, dim):
+            def __init__(self, dim, dtype, use_prim_dtype):
                 super(aten_softmax, self).__init__()
                 self.dim = dim
+                self.dtype = dtype
+                if use_prim_dtype:
+                    self.forward = self.forward_prim_dtype
+                elif dtype is not None:
+                    self.forward = self.forward_dtype
 
             def forward(self, x):
                 return F.softmax(x, self.dim)
 
+            def forward_dtype(self, x):
+                return F.softmax(x, self.dim, dtype=self.dtype)
+
+            def forward_prim_dtype(self, x, y):
+                return F.softmax(x, self.dim, dtype=y.dtype)
+
         ref_net = None
 
-        return aten_softmax(dim), ref_net, "aten::softmax"
+        return aten_softmax(dim, dtype, use_prim_dtype), ref_net, "aten::softmax"
 
     @pytest.mark.parametrize("dim", [-1, 3])
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_softmax(self, dim, ie_device, precision, ir_version):
         self._test(*self.create_model(dim), ie_device, precision, ir_version)
+
+    @pytest.mark.parametrize("dim", [-1, 3])
+    @pytest.mark.parametrize("dtype", ["float32", "float64"])
+    @pytest.mark.parametrize("use_prim_dtype", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_softmax(self, dim, dtype, use_prim_dtype, ie_device, precision, ir_version):
+        input_kwargs = {}
+        if use_prim_dtype:
+            input_kwargs["second_input_dtype"] = dtype
+        self._test(*self.create_model(dim, dtype, use_prim_dtype), ie_device,
+                   precision, ir_version, kwargs_to_prepare_input=input_kwargs)