[PT FE][Opset13] Add translation for torchscript aten::multinomial operator (#20397)

* Enable aten::multinomial in PT FE

* Add aten::multinomial test

* Fix typo

* Update tests/layer_tests/pytorch_tests/test_multinomial.py

Co-authored-by: Roman Kazantsev <roman.kazantsev@intel.com>

* Update multinomial to changes in specification

* rm ;

* Add temporary convert for multinomial i64 support

* Revert convert

* Fix torch test

---------

Co-authored-by: Roman Kazantsev <roman.kazantsev@intel.com>

src/frontends/pytorch/src/op/multinomial.cpp

@@ -0,0 +1,55 @@
+// Copyright (C) 2018-2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/op/multinomial.hpp"
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/concat.hpp"
+#include "openvino/op/gather.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/scatter_elements_update.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_multinomial(const NodeContext& context) {
+    num_inputs_check(context, 3, 5);
+    auto const_neg_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {-1}));
+    auto const_0 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {0}));
+    auto const_1 = context.mark_node(v0::Constant::create(element::i32, Shape{1}, {1}));
+    auto input = context.get_input(0);
+    auto num_samples = context.mark_node(std::make_shared<v1::Reshape>(context.get_input(1), const_1, false));
+    auto replacement = context.const_input<bool>(2);
+    FRONT_END_OP_CONVERSION_CHECK(context.input_is_none(3),
+                                  "aten::multinomial conversion with generator is not supported");
+
+    // Torch multinomial accept input of [class_probs] or [bs, class_probs], convert always to [bs, class_probs] for OV.
+    auto input_shape = context.mark_node(std::make_shared<v3::ShapeOf>(input, element::i32));
+    auto class_probs_shape = context.mark_node(std::make_shared<v8::Gather>(input_shape, const_neg_1, const_0));
+    auto inp_shape = context.mark_node(std::make_shared<v0::Concat>(OutputVector{const_neg_1, class_probs_shape}, 0));
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, inp_shape, false));
+
+    auto multinomial =
+        context.mark_node(std::make_shared<v13::Multinomial>(input, num_samples, element::i64, replacement, false));
+
+    // Torch multinomial can return [num_samples] or [bs, num_samples] based on input dim, reshape to correct shape.
+    auto out_shape = context.mark_node(
+        std::make_shared<v12::ScatterElementsUpdate>(input_shape, const_neg_1, num_samples, const_neg_1));
+    multinomial = context.mark_node(std::make_shared<v1::Reshape>(multinomial, out_shape, false));
+
+    if (!context.input_is_none(5)) {
+        context.mutate_input(5, multinomial);
+    }
+    return {multinomial};
+};
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -126,6 +126,7 @@ OP_CONVERTER(translate_mean);
 OP_CONVERTER(translate_meshgrid);
 OP_CONVERTER(translate_min);
 OP_CONVERTER(translate_minimum);
+OP_CONVERTER(translate_multinomial);
 OP_CONVERTER(translate_narrow);
 OP_CONVERTER(translate_native_multi_head_attention);
 OP_CONVERTER(translate_neg);
@@ -448,6 +449,7 @@ const std::map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::mul_", op::inplace_translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
         {"aten::multiply", op::translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
         {"aten::multiply_", op::inplace_translate_1to1_match_2_inputs_align_types<opset10::Multiply>},
+        {"aten::multinomial", op::translate_multinomial},
         {"aten::narrow", op::translate_narrow},
         {"aten::ne", op::translate_1to1_match_2_inputs_align_types<opset10::NotEqual>},
         {"aten::neg", op::translate_neg},

tests/layer_tests/pytorch_tests/test_multinomial.py

@@ -0,0 +1,99 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestMultinomial(PytorchLayerTest):
+    def _prepare_input(self, input, num_samples, out):
+        model_inputs = [input, np.array(num_samples)]
+        if out:
+            model_inputs.append(np.zeros_like(input).astype(np.int64))
+        return model_inputs
+
+    def create_model(self, replacement, out, test_type):
+        class aten_multinomial(torch.nn.Module):
+            def __init__(self, replacement, out, test_type) -> None:
+                super().__init__()
+                self.replacement = replacement
+                self.test_type = test_type
+                if not out:
+                    self.forward = self.multinomial
+                else:
+                    self.forward = self.multinomial_out
+
+            def multinomial(self, input, num_samples):
+                multinomial = torch.multinomial(input, num_samples, self.replacement)
+                return self.mode(multinomial)
+
+            def multinomial_out(self, input, num_samples, out):
+                multinomial = torch.multinomial(input, num_samples, self.replacement, out=out)
+                return self.mode(multinomial)
+
+            def mode(self, op):
+                if self.test_type == "shape":
+                    return torch.ones_like(op)
+                if self.test_type == "sorted":
+                    return torch.sort(op)[0]
+                return op
+
+        ref_net = None
+
+        return aten_multinomial(replacement, out, test_type), ref_net, "aten::multinomial"
+
+    @pytest.mark.parametrize(
+        ("input", "num_samples", "replacement", "test_type"),
+        [
+            (
+                np.array([[0, 1, 0, 0], [0, 0, 0, 1], [1, 0, 0, 0]], dtype=np.float32),
+                1024,
+                True,
+                "exact",
+            ),
+            (
+                np.array([[0.001, 0.001, 0.1, 0.9], [5, 10, 0.1, 256], [1, 1e-5, 1e-5, 1e-5]], dtype=np.float32),
+                4,
+                False,
+                "sorted",
+            ),
+            (
+                np.array([[0.001, 0, 0.1, 0.9], [5, 10, 0, 256], [0.9, 0.001, 1e-5, 0]], dtype=np.float64),
+                3,
+                False,
+                "sorted",
+            ),
+            (
+                np.array([[1, 2, 3, 4], [4, 3, 2, 1], [1, 0, 0, 0]], dtype=np.float32),
+                1,
+                True,
+                "shape",
+            ),
+            (
+                np.array([1, 2, 3, 4], dtype=np.float32),
+                256,
+                True,
+                "shape",
+            ),
+            (
+                np.array([[1, 2, 3, 4]], dtype=np.float32),
+                256,
+                True,
+                "shape",
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("out", [True, False])
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_multinomial(self, input, num_samples, replacement, out, test_type, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(replacement, out, test_type),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={"input": input, "num_samples": num_samples, "out": out}
+        )