IntenseWebs/openvino
3
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support more complicated cases of list concatenation (#16139)

Browse Source 
* Support more complicated cases of list concatenation

* Fix codestyle

* Add tests
This commit is contained in:
Maxim Vafin authored and GitHub committed 2023-03-10 07:51:10 +01:00
1 parent a1510a5e5f
commit ec8a4abf6d
10 files changed
+213 -104
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/frontends/pytorch/src/op/add.cpp
+1 -1
View File
@@ -24,7 +24,7 @@ OutputVector translate_add(NodeContext& context) {
if (dtype0.is<type::List>() && dtype1.is<type::List>()) {
// aten::add.t(t[] a, t[] b) -> t[]
// Case when two lists gets concatenated
return {context.mark_node(std::make_shared<ov::op::v0::Concat>(OutputVector{lhs, rhs}, 0))};
FRONT_END_OP_CONVERSION_CHECK(false, "aten::add is used for concatenation of lists, not possible to convert");
}
align_eltwise_input_types(context, lhs, rhs);
if (!context.input_is_none(2)) {
src/frontends/pytorch/src/op/cat.cpp
+14 -6
View File
@@ -3,6 +3,7 @@
//
#include "openvino/frontend/pytorch/node_context.hpp"
#include "openvino/op/concat.hpp"
#include "pt_framework_node.hpp"
#include "utils.hpp"
@@ -14,12 +15,19 @@ namespace op {
OutputVector translate_cat(NodeContext& context) {
// This translator is only needed to get axis as constant from external scope
num_inputs_check(context, 2, 2);
auto fw_node = std::make_shared<PtFrameworkNode>(context.get_decoder(), OutputVector{context.get_input(0)}, 1);
auto attrs = fw_node->get_attrs();
// If this fails it means axis is dynamic and aten::cat will be converted to fw node in regular pipeline
attrs["axis"] = std::to_string(context.const_input<int64_t>(1));
fw_node->set_attrs(attrs);
return {context.mark_node(std::dynamic_pointer_cast<Node>(fw_node))};
const auto&& list_elems = get_list_as_outputs(context.get_input(0));
auto axis = context.const_input<int64_t>(1);
if (list_elems.empty()) {
// couldn't get list elements
auto fw_node = std::make_shared<PtFrameworkNode>(context.get_decoder(), OutputVector{context.get_input(0)}, 1);
auto attrs = fw_node->get_attrs();
// If this fails it means axis is dynamic and aten::cat will be converted to fw node in regular pipeline
attrs["axis"] = std::to_string(axis);
fw_node->set_attrs(attrs);
return {context.mark_node(std::dynamic_pointer_cast<Node>(fw_node))};
}
auto concat = std::make_shared<ov::op::v0::Concat>(OutputVector(list_elems.begin(), list_elems.end()), axis);
return {context.mark_node(concat)};
};
} // namespace op
src/frontends/pytorch/src/op/getitem.cpp
+16 -4
View File
@@ -5,6 +5,7 @@
#include "openvino/frontend/pytorch/node_context.hpp"
#include "openvino/op/constant.hpp"
#include "openvino/op/gather.hpp"
#include "pt_framework_node.hpp"
#include "utils.hpp"
namespace ov {
@@ -15,10 +16,21 @@ namespace op {
OutputVector translate_getitem(NodeContext& context) {
num_inputs_check(context, 2, 2);
auto input = context.get_input(0);
FRONT_END_OP_CONVERSION_CHECK(cast_fw_node(input.get_node_shared_ptr(), "prim::ListConstruct") == nullptr,
"unsupported case for aten::getitem");
FRONT_END_OP_CONVERSION_CHECK(cast_fw_node(input.get_node_shared_ptr(), "aten::split") == nullptr,
"unsupported case for aten::getitem");
if (std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(input.get_node_shared_ptr())) {
FRONT_END_OP_CONVERSION_CHECK(!cast_fw_node(input.get_node_shared_ptr(), "aten::split"),
"special case for aten::__getitem__");
const auto&& list_elems = get_list_as_outputs(input);
auto getitem_idx = context.const_input<int64_t>(1);
if (getitem_idx < 0) {
getitem_idx += list_elems.size();
}
FRONT_END_OP_CONVERSION_CHECK(getitem_idx < static_cast<int64_t>(list_elems.size()),
"Index: ",
getitem_idx,
" is out of bounds of input list of len: ",
list_elems.size());
return {list_elems.at(getitem_idx)};
}
auto getitem_idx = context.get_input(1);
auto zero = context.mark_node(ov::op::v0::Constant::create(element::i32, Shape{}, {0}));
return {context.mark_node(std::make_shared<ov::op::v8::Gather>(input, getitem_idx, zero))};
src/frontends/pytorch/src/transforms/aten_cat_replacer.cpp
+4 -19
View File
@@ -88,7 +88,6 @@ AtenCatToConcat::AtenCatToConcat() {
auto list_construct = cast_fw_node(loop->get_input_node_shared_ptr(input_index), "prim::ListConstruct");
if (!list_construct || list_construct->get_input_size() > 0)
return false;
// TODO: Is unsqueeze needed?
auto new_result = std::make_shared<ov::op::v0::Result>(append->input_value(1));
body->add_results({new_result});
auto new_output = loop->get_concatenated_slices(new_result, 0, 1, 1, -1, axis);
@@ -97,25 +96,11 @@ AtenCatToConcat::AtenCatToConcat() {
return true;
}
OutputVector tmp_inputs;
NodeVector rt_copy_from{cat};
while (const auto& input_fw_node = cast_fw_node(input_node, "aten::append")) {
rt_copy_from.push_back(input_fw_node);
tmp_inputs.push_back(input_fw_node->input(1).get_source_output());
input_node = input_fw_node->input(0).get_source_output().get_node_shared_ptr();
}
auto list_construct = cast_fw_node(input_node, "prim::ListConstruct");
if (!list_construct)
return false;
rt_copy_from.push_back(list_construct);
OutputVector inputs;
for (auto& input : list_construct->inputs()) {
inputs.push_back(input.get_source_output());
}
inputs.insert(inputs.end(), tmp_inputs.rbegin(), tmp_inputs.rend());
auto result = std::make_shared<ov::op::v0::Concat>(inputs, axis);
copy_runtime_info(rt_copy_from, result);
const auto&& tmp_inputs = get_list_as_outputs(cat->get_input_source_output(0));
auto result = std::make_shared<ov::op::v0::Concat>(OutputVector(tmp_inputs.begin(), tmp_inputs.end()), axis);
copy_runtime_info(cat, result);
replace_node(cat, result);
result->set_friendly_name(cat->get_friendly_name());
return true;
};
src/frontends/pytorch/src/transforms/listconstruct_replacer.cpp
+35 -27
View File
@@ -16,7 +16,6 @@
#include "openvino/op/shape_of.hpp"
#include "openvino/op/tile.hpp"
#include "openvino/op/transpose.hpp"
#include "openvino/op/unsqueeze.hpp"
#include "openvino/op/util/framework_node.hpp"
#include "openvino/pass/pattern/matcher.hpp"
#include "openvino/pass/pattern/op/or.hpp"
@@ -33,23 +32,23 @@ using namespace ov::op;
ListConstructReplacer::ListConstructReplacer() {
// Transformation for torch operators for cases where prim::ListConstruct can be replaced with Concat.
auto list_construct = pattern::wrap_type<ov::op::util::FrameworkNode>();
auto list = pattern::wrap_type<ov::op::util::FrameworkNode>();
// Both aten::view and aten::reshape are using same translation returning Reshape operator.
auto reshape_op = pattern::wrap_type<v1::Reshape>({pattern::any_input(), list_construct});
auto roll_op = pattern::wrap_type<v7::Roll>({pattern::any_input(), list_construct, pattern::any_input()});
auto broadcast_op = pattern::wrap_type<v3::Broadcast>({pattern::any_input(), list_construct});
auto adapool_op = pattern::wrap_type<v8::AdaptiveAvgPool>({pattern::any_input(), list_construct});
auto reshape_op = pattern::wrap_type<v1::Reshape>({pattern::any_input(), list});
auto roll_op = pattern::wrap_type<v7::Roll>({pattern::any_input(), list, pattern::any_input()});
auto broadcast_op = pattern::wrap_type<v3::Broadcast>({pattern::any_input(), list});
auto adapool_op = pattern::wrap_type<v8::AdaptiveAvgPool>({pattern::any_input(), list});
// replace list construct for aten::expand(tensor, prim::ListConstruct(shapes)) decomposition
// shape_of + broadcast + equal + select
auto shape_of_op = pattern::wrap_type<v3::ShapeOf>({list_construct});
auto equal_op = pattern::wrap_type<v1::Equal>({list_construct, pattern::any_input()});
auto select_op = pattern::wrap_type<v1::Select>({pattern::any_input(), pattern::any_input(), list_construct});
auto shape_of_op = pattern::wrap_type<v3::ShapeOf>({list});
auto equal_op = pattern::wrap_type<v1::Equal>({list, pattern::any_input()});
auto select_op = pattern::wrap_type<v1::Select>({pattern::any_input(), pattern::any_input(), list});
// replace list construct for aten::repeat(tensor, prim::ListConstruct(shapes)))
// shape_of + broadcast + tile
auto tile_op = pattern::wrap_type<v0::Tile>({pattern::any_input(), list_construct});
auto tile_op = pattern::wrap_type<v0::Tile>({pattern::any_input(), list});
// replace aten::permute(tensor, prim::ListConstruct)
auto transpose_op = pattern::wrap_type<v1::Transpose>({pattern::any_input(), list_construct});
auto transpose_op = pattern::wrap_type<v1::Transpose>({pattern::any_input(), list});
auto lc_pattern = std::make_shared<pattern::op::Or>(OutputVector{reshape_op,
roll_op,
broadcast_op,
@@ -63,23 +62,32 @@ ListConstructReplacer::ListConstructReplacer() {
ov::matcher_pass_callback callback = [=](pattern::Matcher& m) {
auto& pattern_map = m.get_pattern_value_map();
auto list_construct_node = pattern_map.at(list_construct).get_node_shared_ptr();
if (auto list_unpack_node = cast_fw_node(list_construct_node, "prim::ListConstruct")) {
// Concatenation is possible because all elements in list should be scalar intigers.
OutputVector inputs;
auto axis_0 = v0::Constant::create(element::i32, Shape{}, {0});
for (auto& input : list_construct_node->inputs()) {
auto rank = input.get_partial_shape().rank();
FRONT_END_OP_CONVERSION_CHECK(rank.is_dynamic() || rank.get_length() == 0, "Rank must be 0");
auto unsqueeze = std::make_shared<v0::Unsqueeze>(input.get_source_output(), axis_0);
inputs.push_back(unsqueeze);
auto list_node = pattern_map.at(list).get_node_shared_ptr();
// Concatenation is possible because all elements in list should be scalar or 1D tensors,
// result should be 1D tensor.
OutputVector inputs;
auto neg_1 = v0::Constant::create(element::i32, Shape{1}, {-1});
const auto& start_output = list_node->output(0);
for (const auto& input : get_list_as_outputs(start_output)) {
if (input == start_output) {
// Start output exist in list elements, it might mean we have only 1 element in list inputs and it is
// already a list, we do not need to concat it
return false;
}
auto concat = std::make_shared<v0::Concat>(inputs, 0);
copy_runtime_info({list_construct_node}, concat);
replace_node(list_construct_node, concat);
return true;
};
return false;
auto rank = input.get_partial_shape().rank();
if (rank.is_static() && rank.get_length() > 1) {
// if list elements of rank higher then 1D we cannot resolve it
return false;
}
// reshape all elements to 1D
auto reshape = std::make_shared<v1::Reshape>(input, neg_1, false);
inputs.push_back(reshape);
}
auto concat = std::make_shared<v0::Concat>(inputs, 0);
copy_runtime_info({list_node}, concat);
replace_node(list_node, concat);
concat->set_friendly_name(list_node->get_friendly_name());
return true;
};
auto m = std::make_shared<pattern::Matcher>(lc_pattern, "ov::frontend::pytorch::pass::ListConstructReplacer");
this->register_matcher(m, callback);
src/frontends/pytorch/src/transforms/prim_list_construct_pad.cpp
+35 -37
View File
@@ -27,30 +27,32 @@ namespace frontend {
namespace pytorch {
namespace pass {
using namespace ov::op;
namespace {
std::shared_ptr<Node> create_padding(std::shared_ptr<Node> input_rank,
std::shared_ptr<Node> padding,
std::shared_ptr<Node> start_id,
std::shared_ptr<Node> end_id) {
Output<Node> create_padding(const Output<Node>& input_rank,
const Output<Node>& padding,
const Output<Node>& start_id,
const Output<Node>& end_id) {
// PyTorch paddings represented as [N_pad_begins, N_pad_ends, N-1_pad_begins, N-1_pad_ends, ... ]
// if len of paddings not equal to input rank * 2, zero padding added to first rank - N dimensions
// OV expects paddings separated on begins and ends for each dimension from first to last
auto minus_two = ov::op::v0::Constant::create(element::i32, Shape{}, {-2});
auto zero = ov::op::v0::Constant::create(element::i32, Shape{}, {0});
auto pad_id_range = std::make_shared<ov::op::v4::Range>(start_id, end_id, minus_two, element::i32);
auto pads = std::make_shared<ov::op::v8::Gather>(padding, pad_id_range, zero);
auto minus_two = v0::Constant::create(element::i32, Shape{}, {-2});
auto zero = v0::Constant::create(element::i32, Shape{}, {0});
auto pad_id_range = std::make_shared<v4::Range>(start_id, end_id, minus_two, element::i32);
auto pads = std::make_shared<v8::Gather>(padding, pad_id_range, zero);
// add left side zero padding for difference between padding size and input rank
auto pads_short_len = std::make_shared<ov::op::v3::ShapeOf>(pads, element::i32);
auto pads_diff = std::make_shared<ov::op::v1::Subtract>(input_rank, pads_short_len);
auto pads_remaining = std::make_shared<ov::op::v3::Broadcast>(zero, pads_diff);
auto pads_remaining_c = std::make_shared<ov::op::v1::ConvertLike>(pads_remaining, pads);
auto pads_full = std::make_shared<ov::op::v0::Concat>(OutputVector{pads_remaining_c, pads}, 0);
auto pads_short_len = std::make_shared<v3::ShapeOf>(pads, element::i32);
auto pads_diff = std::make_shared<v1::Subtract>(input_rank, pads_short_len);
auto pads_remaining = std::make_shared<v3::Broadcast>(zero, pads_diff);
auto pads_remaining_c = std::make_shared<v1::ConvertLike>(pads_remaining, pads);
auto pads_full = std::make_shared<v0::Concat>(OutputVector{pads_remaining_c, pads}, 0);
return pads_full;
}
const std::unordered_map<std::string, ov::op::PadMode> PAD_MODES = {{"constant", ov::op::PadMode::CONSTANT},
{"reflect", ov::op::PadMode::REFLECT},
{"replicate", ov::op::PadMode::EDGE}};
const std::unordered_map<std::string, PadMode> PAD_MODES = {{"constant", PadMode::CONSTANT},
{"reflect", PadMode::REFLECT},
{"replicate", PadMode::EDGE}};
}; // namespace
@@ -62,26 +64,26 @@ PrimListConstructPadReplacer::PrimListConstructPadReplacer() {
if (!pad_op) {
return false;
}
auto minus_two = ov::op::v0::Constant::create(element::i32, Shape{}, {-2});
auto minus_one = ov::op::v0::Constant::create(element::i32, Shape{}, {-1});
auto zero = ov::op::v0::Constant::create(element::i32, Shape{}, {0});
auto input_node = pad_op->input_value(0).get_node_shared_ptr();
auto minus_two = v0::Constant::create(element::i32, Shape{}, {-2});
auto minus_one = v0::Constant::create(element::i32, Shape{}, {-1});
auto zero = v0::Constant::create(element::i32, Shape{}, {0});
auto input_node = pad_op->input_value(0);
auto padding = pad_op->input_value(1).get_node_shared_ptr();
// for case. when padding is list of scalars, concatenate them into one tensor
auto pad_values = concat_list_construct(padding);
std::string mode = "constant";
auto zero_f = ov::op::v0::Constant::create(element::f32, Shape{}, {0});
auto input_shape = std::make_shared<ov::op::v3::ShapeOf>(input_node, element::i32);
auto input_rank = std::make_shared<ov::op::v3::ShapeOf>(input_shape, element::i32);
auto pad_size_1d = std::make_shared<ov::op::v3::ShapeOf>(pad_values, element::i32);
auto pad_size = std::make_shared<ov::op::v0::Squeeze>(pad_size_1d, zero);
auto zero_f = v0::Constant::create(element::f32, Shape{}, {0});
auto input_shape = std::make_shared<v3::ShapeOf>(input_node, element::i32);
auto input_rank = std::make_shared<v3::ShapeOf>(input_shape, element::i32);
auto pad_size_1d = std::make_shared<v3::ShapeOf>(pad_values, element::i32);
auto pad_size = std::make_shared<v0::Squeeze>(pad_size_1d, zero);
// get pad_begins and pad_ends indexes starting for end of paddings
auto start_pad_begins = std::make_shared<ov::op::v1::Add>(pad_size, minus_two);
auto start_pad_ends = std::make_shared<ov::op::v1::Add>(pad_size, minus_one);
auto start_pad_begins = std::make_shared<v1::Add>(pad_size, minus_two);
auto start_pad_ends = std::make_shared<v1::Add>(pad_size, minus_one);
auto pad_begins_full = create_padding(input_rank, pad_values, start_pad_begins, minus_one);
auto pad_ends_full = create_padding(input_rank, pad_values, start_pad_ends, zero);
auto mode_const = pad_op->input_value(2).get_node_shared_ptr();
auto pad_value = pad_op->input_value(3).get_node_shared_ptr();
auto pad_value = pad_op->input_value(3);
if (const auto& fw_node_mode = cast_fw_node(mode_const, "prim::Constant")) {
const auto& attrs = fw_node_mode->get_attrs();
if (attrs.find("string_value") != attrs.end()) {
@@ -89,27 +91,23 @@ PrimListConstructPadReplacer::PrimListConstructPadReplacer() {
}
}
if (mode == "constant") {
if (const auto& fw_node_value = cast_fw_node(pad_value, "prim::Constant")) {
if (const auto& fw_node_value = cast_fw_node(pad_value.get_node_shared_ptr(), "prim::Constant")) {
const auto& attrs = fw_node_value->get_attrs();
if (attrs.find("none_value") != attrs.end()) {
pad_value = zero_f;
}
}
pad_value = std::make_shared<ov::op::v1::ConvertLike>(pad_value, input_node);
pad_value = std::make_shared<v1::ConvertLike>(pad_value, input_node);
}
FRONT_END_OP_CONVERSION_CHECK(PAD_MODES.find(mode) != PAD_MODES.end(),
"Unsupported mode: ",
mode,
"for aten::pad");
auto pad_mode = PAD_MODES.at(mode);
auto pad = std::make_shared<ov::op::v1::Pad>(input_node, pad_begins_full, pad_ends_full, pad_value, pad_mode);
auto pad = std::make_shared<v1::Pad>(input_node, pad_begins_full, pad_ends_full, pad_value, pad_mode);
replace_node(pad_op, pad);
copy_runtime_info({pad_op,
input_node,
padding,
pad_op->input_value(2).get_node_shared_ptr(),
pad_op->input_value(3).get_node_shared_ptr()},
pad);
copy_runtime_info({pad_op, padding, mode_const, pad_op->input_value(3).get_node_shared_ptr()}, pad);
pad->set_friendly_name(pad_op->get_friendly_name());
return true;
};
src/frontends/pytorch/src/utils.cpp
+31
View File
@@ -393,6 +393,37 @@ void align_eltwise_input_types(const NodeContext& context, Output<Node>& lhs, Ou
}
}
std::deque<Output<Node>> get_list_as_outputs(const Output<Node>& start) {
std::deque<Output<Node>> res;
auto current_output = start;
while (const auto& input_fw_node =
std::dynamic_pointer_cast<ov::op::util::FrameworkNode>(current_output.get_node_shared_ptr())) {
const auto& attrs = input_fw_node->get_attrs();
if (attrs.find("PtTypeName") == attrs.end()) {
break;
}
if (attrs.at("PtTypeName") == "aten::append") {
res.push_front(input_fw_node->input(1).get_source_output());
} else if (attrs.at("PtTypeName") == "aten::add") {
const auto&& lhs_list = get_list_as_outputs(input_fw_node->input(1).get_source_output());
res.insert(res.end(), lhs_list.begin(), lhs_list.end());
} else {
break;
}
current_output = input_fw_node->input(0).get_source_output();
}
auto list_construct = cast_fw_node(current_output.get_node_shared_ptr(), "prim::ListConstruct");
if (list_construct) {
auto inputs = list_construct->inputs();
for (auto input_it = inputs.rbegin(); input_it != inputs.rend(); ++input_it) {
res.push_front(input_it->get_source_output());
}
} else {
res.push_front(current_output);
}
return res;
}
} // namespace pytorch
} // namespace frontend
} // namespace ov
src/frontends/pytorch/src/utils.hpp
+3 -1
View File
@@ -49,7 +49,7 @@ OutputVector make_framework_node(NodeContext& context);
std::shared_ptr<op::util::FrameworkNode> cast_fw_node(std::shared_ptr<Node> node, const std::string& type);
// TODO: Elimitate the need of this function by implementing more accurate custom data type handling
// TODO: Eliminate the need of this function by implementing more accurate custom data type handling
Any simplified_type_interpret(Any type);
void align_eltwise_input_types(const NodeContext& context,
@@ -57,6 +57,8 @@ void align_eltwise_input_types(const NodeContext& context,
Output<Node>& rhs,
bool align_scalars = false);
std::deque<Output<Node>> get_list_as_outputs(const Output<Node>& start);
namespace op {
template <OutputVector (*T)(NodeContext&), size_t idx = 0>
OutputVector inplace_op(NodeContext& context) {
tests/layer_tests/pytorch_tests/test_cat.py
+14
View File
@@ -19,6 +19,7 @@ class aten_append_cat(torch.nn.Module):
list.append(x)
return torch.cat(list, 1)
class aten_loop_append_cat(torch.nn.Module):
def forward(self, x):
list = []
@@ -27,6 +28,13 @@ class aten_loop_append_cat(torch.nn.Module):
return torch.cat(list, 1)
class aten_add_cat(torch.nn.Module):
def forward(self, x):
list = [x, x]
list2 = list + [x, x]
return torch.cat(list2, 1)
class TestCat(PytorchLayerTest):
def _prepare_input(self):
import numpy as np
@@ -49,3 +57,9 @@ class TestCat(PytorchLayerTest):
def test_loop_append_cat(self, ie_device, precision, ir_version):
self._test(aten_loop_append_cat(), None, ["aten::cat", "aten::append", "prim::ListConstruct", "prim::Loop"],
ie_device, precision, ir_version, freeze_model=False)
@pytest.mark.nightly
@pytest.mark.precommit
def test_add_cat(self, ie_device, precision, ir_version):
self._test(aten_add_cat(), None, ["aten::cat", "aten::add", "prim::ListConstruct"],
ie_device, precision, ir_version, freeze_model=False)
tests/layer_tests/pytorch_tests/test_getitem.py
+60 -9
View File
@@ -2,6 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
import pytest
import torch
from pytorch_layer_test_class import PytorchLayerTest
@@ -12,19 +13,19 @@ class TestGetItem(PytorchLayerTest):
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
self.idx: int = idx
def forward(self, x):
if x.shape[self.idx] > self.idx:
res = x.shape[self.idx]
@@ -44,11 +45,61 @@ class TestGetItem(PytorchLayerTest):
@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)])
([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})
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)