[PT FE] Support aten::take operation

src/frontends/pytorch/src/op/take.cpp

@@ -0,0 +1,43 @@
+// Copyright (C) 2018-2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/core/validation_util.hpp"
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/constant.hpp"
+#include "openvino/op/gather.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/shape_of.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_take_op(const NodeContext& context) {
+    num_inputs_check(context, 2, 3);
+    auto input = context.get_input(0);
+    auto indices = context.get_input(1);
+    auto input_shape = input.get_partial_shape();
+    if (input_shape.rank().is_static() && input_shape.rank().get_length() == 0) {
+        FRONT_END_OP_CONVERSION_CHECK(false, "input tensor MUST be non-scalar");
+    }
+    auto new_shape = context.mark_node(v0::Constant::create(element::i64, Shape{1}, {-1}));
+    input = context.mark_node(std::make_shared<v1::Reshape>(input, new_shape, false));
+    indices = context.mark_node(std::make_shared<v0::Convert>(indices, element::i64));
+    auto axis_constant = context.mark_node(v0::Constant::create(element::i64, Shape{}, {0}));
+    auto gather = context.mark_node(std::make_shared<v8::Gather>(input, indices, axis_constant));
+    if (!context.input_is_none(2)) {
+        context.mutate_input(2, gather);
+    }
+
+    return {gather};
+}
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -255,6 +255,7 @@ OP_CONVERTER(translate_sub_);
 OP_CONVERTER(translate_sum);
 OP_CONVERTER(translate_t);
 OP_CONVERTER(translate_take_along_dim);
+OP_CONVERTER(translate_take_op);
 OP_CONVERTER(translate_to);
 OP_CONVERTER(translate_topk);
 OP_CONVERTER(translate_transpose);
@@ -725,6 +726,7 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
         {"aten::swapaxes", op::quantizable_op<op::translate_transpose>},
         {"aten::t", op::translate_t},
         {"aten::take_along_dim", op::translate_take_along_dim},
+        {"aten::take", op::translate_take_op},
         {"aten::tan", op::optional_out<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Tan>, 1>},
         {"aten::tan_", op::inplace_op<op::translate_1to1_match_1_inputs<opset10::Tan>>},
         {"aten::tanh", op::optional_out<op::translate_1to1_match_1_inputs_with_fp32_type_alignment<opset10::Tanh>, 1>},

tests/layer_tests/pytorch_tests/test_take.py

@@ -0,0 +1,35 @@
+# Copyright (C) 2018-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import numpy as np
+import torch
+from pytorch_layer_test_class import PytorchLayerTest
+
+class TestTake(PytorchLayerTest):
+    def _prepare_input(self, input_shape, indices_shape, max_val):
+        input_tensor = np.random.randn(*input_shape).astype(np.float32)
+        indices = np.random.randint(-max_val, max_val, indices_shape).astype(np.int64)
+        return (input_tensor, indices)
+
+    def create_model(self):
+        class aten_take(torch.nn.Module):
+            def forward(self, x, indices):
+                return torch.take(x, indices)
+
+        ref_net = None
+        return aten_take(), ref_net, "aten::take"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.parametrize("input_shape", [(10,), (3, 4), (2, 3, 4)])
+    @pytest.mark.parametrize("indices_shape", [(5,), (2, 2), (3, 2)])
+    def test_take(self, input_shape, indices_shape, ie_device, precision, ir_version):
+        max_val = np.prod(input_shape)
+        self._test(*self.create_model(), ie_device, precision, ir_version, 
+                   kwargs_to_prepare_input={
+                       "input_shape": input_shape,
+                       "indices_shape": indices_shape,
+                       "max_val": max_val
+                   })