54 lines
2.3 KiB
Python
54 lines
2.3 KiB
Python
# Copyright (C) 2018-2023 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pytest
|
|
|
|
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"):
|
|
import torch
|
|
|
|
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}) |