[PT FE] Added support for aten::__lshift__ and aten::__rshift__ and aten::bitwise_left_shift and aten::bitwise_right_shift

src/frontends/pytorch/src/op/lshift.cpp

@@ -0,0 +1,29 @@
+// Copyright (C) 2018-2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/op/bitwise_left_shift.hpp"
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "utils.hpp"                     
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_lshift(const NodeContext& context) {
+    num_inputs_check(context, 2, 2);
+
+    auto [input_tensor, shift_amount] = get_inputs_with_promoted_types(context, 0, 1);
+
+    auto lshift_node = context.mark_node(std::make_shared<v15::BitwiseLeftShift>(input_tensor, shift_amount));
+
+    return {lshift_node};
+}
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op/rshift.cpp

@@ -0,0 +1,29 @@
+// Copyright (C) 2018-2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "openvino/frontend/pytorch/node_context.hpp"
+#include "openvino/op/bitwise_right_shift.hpp"
+#include "utils.hpp"
+
+namespace ov {
+namespace frontend {
+namespace pytorch {
+namespace op {
+
+using namespace ov::op;
+
+OutputVector translate_rshift(const NodeContext& context) {
+    num_inputs_check(context, 2, 2);
+
+    auto [input_tensor, shift_amount] = get_inputs_with_promoted_types(context, 0, 1);
+
+    auto rshift_node = context.mark_node(std::make_shared<v15::BitwiseRightShift>(input_tensor, shift_amount));
+
+    return {rshift_node};
+}
+
+}  // namespace op
+}  // namespace pytorch
+}  // namespace frontend
+}  // namespace ov

src/frontends/pytorch/src/op_table.cpp

@@ -52,6 +52,8 @@ OP_CONVERTER(translate_bernoulli);
 OP_CONVERTER(translate_bitwise_and);
 OP_CONVERTER(translate_bitwise_not);
 OP_CONVERTER(translate_bitwise_or);
+OP_CONVERTER(translate_lshift);
+OP_CONVERTER(translate_rshift);
 OP_CONVERTER(translate_bitwise_xor);
 OP_CONVERTER(translate_bucketize);
 OP_CONVERTER(translate_cat);
@@ -343,6 +345,10 @@ const std::unordered_map<std::string, CreatorFunction> get_supported_ops_ts() {
     return {
         {"aten::__and__", op::translate_bitwise_and},
         {"aten::__iand__", op::inplace_op<op::translate_bitwise_and>},
+        {"aten::__lshift__", op::translate_lshift},
+        {"aten::__rshift__", op::translate_rshift},
+        {"aten::bitwise_left_shift", op::translate_lshift},
+        {"aten::bitwise_right_shift", op::translate_rshift},
         {"aten::__derive_index", op::translate_derive_index},
         {"aten::__getitem__", op::translate_getitem},
         {"aten::__not__", op::translate_1to1_match_1_inputs<opset10::LogicalNot>},

tests/layer_tests/pytorch_tests/test_shift_operations.py

@@ -0,0 +1,108 @@
+# Copyright (C) 2018-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+from pytorch_layer_test_class import PytorchLayerTest, skip_if_export
+
+
+class TestShiftOperators(PytorchLayerTest):
+    def _prepare_input(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape):
+        choices = np.array([1, 2, 4, 8, 16, 32])  
+        shifts = np.array([0, 1, 2, 3, 4, 5])     
+
+        x = np.random.choice(choices, lhs_shape).astype(lhs_dtype)
+        y = np.random.choice(shifts, rhs_shape).astype(rhs_dtype)
+        return x, y
+
+    def create_model(self):
+        class aten_shift(torch.nn.Module):
+            def forward(self, lhs, rhs):
+                return lhs << rhs, lhs >> rhs 
+
+        ref_net = None
+        return aten_shift(), ref_net, ("aten::__lshift__", "aten::__rshift__")
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    @pytest.mark.parametrize("lhs_dtype", ["int32", "int64"])
+    @pytest.mark.parametrize("rhs_dtype", ["int32", "int64"])
+    @pytest.mark.parametrize(
+        ("lhs_shape", "rhs_shape"),
+        [
+            ([2, 3], [2, 3]),  
+            ([2, 3], []),      
+            ([], [2, 3]),  
+            ([], []), 
+        ],
+    )
+    def test_shift_operators(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={
+                "lhs_dtype": lhs_dtype,
+                "rhs_dtype": rhs_dtype,
+                "lhs_shape": lhs_shape,
+                "rhs_shape": rhs_shape,
+            },
+            trace_model=True,
+            freeze_model=False,
+        )
+
+
+class TestBitwiseShiftFunctions(PytorchLayerTest):
+    def _prepare_input(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape):
+        choices = np.array([1, 2, 4, 8, 16, 32])  
+        shifts = np.array([0, 1, 2, 3, 4, 5])    
+
+        x = np.random.choice(choices, lhs_shape).astype(lhs_dtype)
+        y = np.random.choice(shifts, rhs_shape).astype(rhs_dtype)
+        return x, y
+
+    def create_model(self):
+        class aten_bitwise_shift(torch.nn.Module):
+            def forward(self, lhs, rhs):
+                return (
+                    torch.bitwise_left_shift(lhs, rhs),
+                    torch.bitwise_right_shift(lhs, rhs)
+                )
+
+        ref_net = None
+        return aten_bitwise_shift(), ref_net, ("aten::bitwise_left_shift", "aten::bitwise_right_shift")
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.precommit_torch_export
+    @pytest.mark.precommit_fx_backend
+    @pytest.mark.parametrize("lhs_dtype", ["int32", "int64"])
+    @pytest.mark.parametrize("rhs_dtype", ["int32", "int64"])
+    @pytest.mark.parametrize(
+        ("lhs_shape", "rhs_shape"),
+        [
+            ([2, 3], [2, 3]),  
+            ([2, 3], []),      
+            ([], [2, 3]),      
+            ([], []), 
+        ],
+    )
+    def test_bitwise_shift_functions(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={
+                "lhs_dtype": lhs_dtype,
+                "rhs_dtype": rhs_dtype,
+                "lhs_shape": lhs_shape,
+                "rhs_shape": rhs_shape,
+            },
+            trace_model=True,
+            freeze_model=False,
+        )