[PT FE]: support aten::masked_scatter (#19038)

src/frontends/pytorch/src/op/masked_scatter.cpp

@@ -0,0 +1,60 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/broadcast.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/gather.hpp"
+#include "openvino/op/non_zero.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/scatter_nd_update.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "openvino/op/slice.hpp"
+#include "openvino/op/transpose.hpp"
+#include "openvino/op/unsqueeze.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_masked_scatter(const NodeContext& context) {
+    // aten::masked_scatter(Tensor self, Tensor mask, Tensor source) -> Tensor
+    // aten::masked_scatter.out(Tensor self, Tensor mask, Tensor source, *, Tensor(a!) out) -> Tensor(a!)
+    // aten::masked_scatter_(Tensor(a!) self, Tensor mask, Tensor source) -> Tensor(a!)
+    num_inputs_check(context, 3, 4);
+    auto x = context.get_input(0);
+    auto mask = context.get_input(1);
+    auto source = context.get_input(2);
+    // mask should be broadcastable to x shape
+    auto x_shape = context.mark_node(std::make_shared<v3::ShapeOf>(x, element::i32));
+    auto expanded_mask =
+        context.mark_node(std::make_shared<v3::Broadcast>(mask, x_shape, BroadcastType::BIDIRECTIONAL));
+    auto index = context.mark_node(std::make_shared<v3::NonZero>(expanded_mask, element::i32));
+    auto input_order = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {1, 0}));
+    index = context.mark_node(std::make_shared<v1::Transpose>(index, input_order));
+    // source can be arbitary shape, select only relevant data
+    auto const_minus_one = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1}));
+    auto flatten_source = context.mark_node(std::make_shared<v1::Reshape>(source, const_minus_one, false));
+    auto const_zero = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0}));
+    auto index_shape = context.mark_node(std::make_shared<v3::ShapeOf>(index, element::i32));
+    auto index_dim_zero = context.mark_node(std::make_shared<v8::Gather>(index_shape, const_zero, const_zero));
+    auto slice_steps = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1}));
+    auto sliced_source = context.mark_node(
+        std::make_shared<v8::Slice>(flatten_source, const_zero, index_dim_zero, slice_steps, const_zero));
+
+    auto update = context.mark_node(std::make_shared<v3::ScatterNDUpdate>(x, index, sliced_source));
+    if (!context.input_is_none(3)) {
+        context.mutate_input(3, update);
+    }
+    return {update};
+};
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -93,6 +93,7 @@ OP_CONVERTER(translate_log2);
 OP_CONVERTER(translate_logsumexp);
 OP_CONVERTER(translate_loop);
 OP_CONVERTER(translate_masked_fill);
+OP_CONVERTER(translate_masked_scatter);
 OP_CONVERTER(translate_max);
 OP_CONVERTER(translate_max_poolnd);
 OP_CONVERTER(translate_mean);
@@ -341,6 +342,8 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::lt", op::translate_1to1_match_2_inputs_align_types<opset10::Less>},
         {"aten::masked_fill", op::translate_masked_fill},
         {"aten::masked_fill_", op::inplace_op<op::translate_masked_fill>},
+        {"aten::masked_scatter", op::translate_masked_scatter},
+        {"aten::masked_scatter_", op::inplace_op<op::translate_masked_scatter>},
         {"aten::matmul", op::translate_1to1_match_2_inputs<opset10::MatMul>},
         {"aten::max", op::translate_max},
         {"aten::max_pool1d", op::quantizable_op<op::translate_max_poolnd>},

tests/layer_tests/pytorch_tests/test_masked_scatter.py

@@ -0,0 +1,61 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestMaskedScatter(PytorchLayerTest):
+    def _prepare_input(self, shape, x_dtype="float32", mask_dtype="bool", out=False):
+        import numpy as np
+        x = np.random.randn(*shape).astype(x_dtype)
+        mask = (x > 0.5).astype(mask_dtype)
+        source = np.arange(np.size(x)).reshape(shape).astype(x_dtype)
+        if not out:
+            return (x, mask, source)
+        y = np.zeros_like(x).astype(x_dtype)
+        return (x, mask, source, y)
+
+    def create_model(self, out=False, inplace=False):
+        import torch
+
+        class aten_masked_scatter(torch.nn.Module):
+            def __init__(self, out, inplace):
+                super(aten_masked_scatter, self).__init__()
+                if inplace:
+                    self.forward = self.forward_inplace
+                if out:
+                    self.forward = self.forward_out
+
+            def forward(self, x, mask, source):
+                return torch.masked_scatter(x, mask, source)
+            
+            def forward_out(self, x, mask, source, out):
+                return torch.masked_scatter(x, mask, source, out=out), out
+
+            def forward_inplace(self, x, mask, source):
+                return x.masked_scatter_(mask, source), x
+
+        ref_net = None
+
+        return aten_masked_scatter(out, inplace), ref_net, "aten::masked_scatter" if not inplace else "aten::masked_scatter_"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("shape", [[2, 5], [10, 10], [2, 3, 4], [10, 5, 10, 3], [2, 6, 4, 1]])
+    @pytest.mark.parametrize("input_dtype", ["float32", "int32", "float", "int", "uint8"])
+    @pytest.mark.parametrize("mask_dtype", ["bool", "uint8"])
+    @pytest.mark.parametrize("out", [True, False])
+    def test_masked_scatter(self, shape, input_dtype, mask_dtype, out, ie_device, precision, ir_version):
+        self._test(*self.create_model(out), ie_device, precision, ir_version, 
+                   kwargs_to_prepare_input={"shape": shape, "x_dtype": input_dtype, "mask_dtype": mask_dtype, "out": out})
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("shape", [[2, 5], [10, 10], [2, 3, 4], [10, 5, 10, 3], [2, 6, 4, 1]])
+    @pytest.mark.parametrize("input_dtype", ["float32", "int32", "float", "int", "uint8"])
+    @pytest.mark.parametrize("mask_dtype", ["bool", "uint8"])
+    def test_masked_scatter_inplace(self, shape, input_dtype, mask_dtype, ie_device, precision, ir_version):
+        self._test(*self.create_model(inplace=True), ie_device, precision, ir_version, 
+                   kwargs_to_prepare_input={"shape": shape, "x_dtype": input_dtype, "mask_dtype": mask_dtype})