From 8c5bf41b14b83a07c21d1b7790eed2f2db0d2639 Mon Sep 17 00:00:00 2001 From: Maxim Vafin Date: Thu, 26 Oct 2023 18:26:21 +0200 Subject: [PATCH] Add test --- tests/layer_tests/pytorch_tests/test_loop.py | 45 ++++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 tests/layer_tests/pytorch_tests/test_loop.py diff --git a/tests/layer_tests/pytorch_tests/test_loop.py b/tests/layer_tests/pytorch_tests/test_loop.py new file mode 100644 index 00000000000000..8c91833f437b18 --- /dev/null +++ b/tests/layer_tests/pytorch_tests/test_loop.py @@ -0,0 +1,45 @@ +# Copyright (C) 2018-2023 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +import os + +import pytest +import numpy as np + +from pytorch_layer_test_class import PytorchLayerTest + + +class TestLoopWithAlias(PytorchLayerTest): + def _prepare_input(self): + return (np.random.randn(*self.shape).astype(np.float32),) + + def create_model(self, n): + import torch + + class loop_alias_model(torch.nn.Module): + def __init__(self, n): + super(loop_alias_model, self).__init__() + self.n = n + + def forward(self, x): + N = x.shape[1] + res = torch.zeros(1, self.n, dtype=torch.long) + d = torch.ones(1, N) * 1e10 + f = torch.zeros(1, dtype=torch.long) + + for i in range(self.n): + res[:, i] = f + _d = torch.sum((x - x[0, f, :]) ** 2, -1) + m = _d < d + d[m] = _d[m] + f = torch.max(d, -1)[1] + return res + + return loop_alias_model(n), None, ["prim::Loop", "aten::copy_"] + + @pytest.mark.parametrize("s,n", [([1, 1024, 3], 512), ([1, 512, 3], 128)]) + @pytest.mark.nightly + @pytest.mark.precommit + def test_loop_alias(self, s, n, ie_device, precision, ir_version): + self.shape = s + self._test(*self.create_model(n), ie_device, precision, + ir_version, use_convert_model=True)