106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
# Copyright (C) 2018-2023 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from pytorch_layer_test_class import PytorchLayerTest
|
|
|
|
|
|
class TestGetItem(PytorchLayerTest):
|
|
def _prepare_input(self, input_shape):
|
|
import numpy as np
|
|
return (np.random.randn(*input_shape).astype(np.float32),)
|
|
|
|
def create_model(self, idx, case="size_with_getitem"):
|
|
class aten_size_get_item(torch.nn.Module):
|
|
def __init__(self, idx):
|
|
super().__init__()
|
|
self.idx = idx
|
|
|
|
def forward(self, x):
|
|
return x.shape[self.idx]
|
|
|
|
class aten_size_get_item_with_if(torch.nn.Module):
|
|
def __init__(self, idx):
|
|
super().__init__()
|
|
self.idx: int = idx
|
|
|
|
def forward(self, x):
|
|
if x.shape[self.idx] > self.idx:
|
|
res = x.shape[self.idx]
|
|
else:
|
|
res = x.shape[-self.idx]
|
|
return res
|
|
|
|
ref_net = None
|
|
op_cls = {
|
|
"getitem": (aten_size_get_item, ["aten::size", "aten::__getitem__"]),
|
|
"getitem_with_if": (aten_size_get_item_with_if, ["aten::size", "aten::__getitem__", "prim::If"])
|
|
}
|
|
op, op_in_graph = op_cls[case]
|
|
|
|
return op(idx), ref_net, op_in_graph
|
|
|
|
@pytest.mark.nightly
|
|
@pytest.mark.precommit
|
|
@pytest.mark.parametrize(("input_shape", "idx"), [
|
|
([1,], 0),
|
|
([1,], -1),
|
|
([1, 2], 0),
|
|
([1, 2], 1),
|
|
([1, 2], -1),
|
|
([1, 2], -2),
|
|
([1, 2, 3], 0),
|
|
([1, 2, 3], 1),
|
|
([1, 2, 3], 2),
|
|
([1, 2, 3], -1),
|
|
([1, 2, 3], -2),
|
|
([1, 2, 3], -3),
|
|
([1, 2, 3, 4], 0),
|
|
([1, 2, 3, 4], 1),
|
|
([1, 2, 3, 4], 2),
|
|
([1, 2, 3, 4], 3),
|
|
([1, 2, 3, 4], -1),
|
|
([1, 2, 3, 4], -2),
|
|
([1, 2, 3, 4], -3),
|
|
([1, 2, 3, 4], -4),
|
|
([1, 2, 3, 4, 5], 0),
|
|
([1, 2, 3, 4, 5], 1),
|
|
([1, 2, 3, 4, 5], 2),
|
|
([1, 2, 3, 4, 5], 3),
|
|
([1, 2, 3, 4, 5], 4),
|
|
([1, 2, 3, 4, 5], -1),
|
|
([1, 2, 3, 4, 5], -2),
|
|
([1, 2, 3, 4, 5], -3),
|
|
([1, 2, 3, 4, 5], -4),
|
|
([1, 2, 3, 4, 5], -5)])
|
|
@pytest.mark.parametrize("case", ["getitem", "getitem_with_if"])
|
|
def test_getitem(self, input_shape, idx, case, ie_device, precision, ir_version):
|
|
self._test(*self.create_model(idx, case), ie_device, precision,
|
|
ir_version, kwargs_to_prepare_input={"input_shape": input_shape})
|
|
|
|
|
|
class aten_add_getitem(torch.nn.Module):
|
|
def __init__(self, idx):
|
|
super().__init__()
|
|
self.idx = idx
|
|
|
|
def forward(self, x):
|
|
list = [x, 2*x]
|
|
list2 = list + [3*x, 4*x]
|
|
return list2[self.idx]
|
|
|
|
|
|
class TestAddGetItem(PytorchLayerTest):
|
|
def _prepare_input(self):
|
|
import numpy as np
|
|
return (np.random.randn(2, 1, 3),)
|
|
|
|
@pytest.mark.nightly
|
|
@pytest.mark.precommit
|
|
@pytest.mark.parametrize("idx", [-4, -3, -2, -1, 0, 1, 2, 3])
|
|
def test_add_cat(self, ie_device, precision, ir_version, idx):
|
|
self._test(aten_add_getitem(idx), None, ["aten::__getitem__", "aten::add", "prim::ListConstruct"],
|
|
ie_device, precision, ir_version, freeze_model=False, use_convert_model=True)
|