From 147d0af7bea9161cb5be9d9c3d23d70be06b54f5 Mon Sep 17 00:00:00 2001 From: Daniil Lyakhov Date: Mon, 25 Nov 2024 13:07:31 -0800 Subject: [PATCH] [Frontends][TorchFX] torch.ops.aten._unsafe_index support (#27617) ### Details: - torch.ops.aten._unsafe_index is supported by the Torch frontend - To support Yolo(V8 and V11) TorchFX models quantization with NNCF ### Tickets: - 141640 --------- Co-authored-by: Roman Kazantsev --- .../frontend/pytorch/torchdynamo/op_support.py | 1 + src/frontends/pytorch/src/op_table.cpp | 1 + tests/layer_tests/pytorch_tests/test_index_tensor.py | 11 +++++++---- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py index f1d0fdb0d36955..b86ae857847a1f 100644 --- a/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py +++ b/src/bindings/python/src/openvino/frontend/pytorch/torchdynamo/op_support.py @@ -145,6 +145,7 @@ def __init__(self, options): "torch.ops.aten.hardtanh.default": None, "torch.ops.aten.hardtanh_.default": None, "torch.ops.aten.index.Tensor": None, + "torch.ops.aten._unsafe_index.Tensor": None, "torch.ops.aten.index_select.default": None, "torch.ops.aten.isfinite.default": None, "torch.ops.aten.isinf.default": None, diff --git a/src/frontends/pytorch/src/op_table.cpp b/src/frontends/pytorch/src/op_table.cpp index 607f0bd32db80d..0fd2d8e54006a4 100644 --- a/src/frontends/pytorch/src/op_table.cpp +++ b/src/frontends/pytorch/src/op_table.cpp @@ -862,6 +862,7 @@ const std::unordered_map get_supported_ops_fx() { {"aten.hardtanh.default", op::translate_hardtanh}, {"aten.hardtanh_.default", op::inplace_op}, {"aten.index.Tensor", op::translate_index_fx}, + {"aten._unsafe_index.Tensor", op::translate_index_fx}, {"aten.index_select.default", op::translate_index_select}, {"aten.isfinite.default", op::translate_1to1_match_1_inputs}, {"aten.isinf.default", op::translate_1to1_match_1_inputs}, diff --git a/tests/layer_tests/pytorch_tests/test_index_tensor.py b/tests/layer_tests/pytorch_tests/test_index_tensor.py index d2055b5f5a4ec5..13782965b4dc26 100644 --- a/tests/layer_tests/pytorch_tests/test_index_tensor.py +++ b/tests/layer_tests/pytorch_tests/test_index_tensor.py @@ -11,7 +11,7 @@ def _prepare_input(self, input_shape): import numpy as np return (np.random.randn(*input_shape).astype(np.float32),) - def create_model(self, indices_list): + def create_model(self, indices_list, safe: bool): import torch class aten_index_tensor(torch.nn.Module): @@ -20,7 +20,9 @@ def __init__(self, indices_list): self.indices_list = indices_list def forward(self, x): - return torch.ops.aten.index.Tensor(x, self.indices_list) + if safe: + return torch.ops.aten.index.Tensor(x, self.indices_list) + return torch.ops.aten._unsafe_index.Tensor(x, self.indices_list) ref_net = None @@ -35,6 +37,7 @@ def forward(self, x): @pytest.mark.nightly @pytest.mark.precommit_torch_export + @pytest.mark.parametrize('safe', [True, False]) @pytest.mark.parametrize(('input_shape', 'indices_list'), [ ([3, 7], [[0], [5, 3, 0]]), ([3, 7, 6], [[0], None, None]), @@ -42,8 +45,8 @@ def forward(self, x): ([3, 7, 6], [[0, 2, 1], None, [5, 0, 3]]), ([3, 7, 6], [[0, 2, 1], [4], [5, 0, 3]]), ]) - def test_index_tensor(self, input_shape, indices_list, ie_device, precision, ir_version): + def test_index_tensor(self, safe, input_shape, indices_list, ie_device, precision, ir_version): if not PytorchLayerTest.use_torch_export(): pytest.skip(reason='aten.index.Tensor test is supported only on torch.export()') - self._test(*self.create_model(indices_list), ie_device, precision, ir_version, + self._test(*self.create_model(indices_list, safe), ie_device, precision, ir_version, kwargs_to_prepare_input={'input_shape': input_shape})