fix TSGatherForward transformation (#19821)

* fix ts_gather_forward

* remove unneeded header
This commit is contained in:
Evgeny Kotov 2023-09-19 11:55:43 +02:00 committed by GitHub
parent 9c908f5245
commit 35d0d92ef7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 44 additions and 7 deletions

View File

@ -57,16 +57,17 @@ TSGatherForward::TSGatherForward() {
}
}
size_t axis;
size_t order_axis;
if (axes[0] < 0) {
auto data_rank = main_node->get_input_partial_shape(0).rank();
if (data_rank.is_dynamic()) {
return false;
}
axis = static_cast<size_t>(axes[0] + data_rank.get_length());
order_axis = static_cast<size_t>(axes[0] + data_rank.get_length());
} else {
axis = static_cast<size_t>(axes[0]);
order_axis = static_cast<size_t>(axes[0]);
}
const size_t axis = order_val[order_axis];
/*
https://docs.openvino.ai/2023.0/openvino_docs_ops_movement_Gather_8.html
The Gather output shape has the same shape as the input,
@ -136,8 +137,7 @@ TSGatherForward::TSGatherForward() {
if (!success) {
return false;
}
auto new_axis =
ov::op::v0::Constant::create(gather_axis->get_element_type(), gather_axis->get_shape(), {order_val[axis]});
auto new_axis = ov::op::v0::Constant::create(gather_axis->get_element_type(), gather_axis->get_shape(), {axis});
main_node->input(2).replace_source_output(new_axis);
copy_runtime_info(gather_axis, new_axis);

View File

@ -63,6 +63,7 @@ auto wrapper = [](const TestCase& test_case) {
struct GatherForwardArguments {
OutputVector inputs_to_main;
std::function<OutputVector(const vector<size_t>&, const OutputVector&)> create_input_transpose_to_main;
Output<Node> new_Gather_first_input;
AxisVector new_transpose_order;
};
@ -76,7 +77,7 @@ auto test_forward_gather = [](const GatherForwardArguments& test_arguments) {
test_case.inputs_to_main = test_arguments.inputs_to_main;
// Test model description:
test_case.model.preprocess_inputs_to_main = {{set_transpose_for}, {{0}}};
test_case.model.preprocess_inputs_to_main = {{test_arguments.create_input_transpose_to_main}, {{0}}};
test_case.model.main_op = {CREATE_GATHER_FACTORY(Gather)};
test_case.model.model_template = create_model;
@ -104,24 +105,44 @@ auto test_forward_gather = [](const GatherForwardArguments& test_arguments) {
return wrapper(test_case);
};
class SetTransposeWithOrder {
public:
SetTransposeWithOrder(const AxisVector& order) : _order(order) {}
OutputVector operator()(const vector<size_t>& idxs, const OutputVector& out_vec) const {
return set_transpose_with_order(idxs, out_vec, _order);
}
private:
const AxisVector _order;
};
vector<GatherForwardArguments> tests_arguments_fw{
{{{parameter(f32, {3, 4, 5, 6}), constant<int>(i32, {2}, {0, 2}), constant<int>(i32, {1}, {2})}},
set_transpose_for,
constant<int>(i32, {1}, {1}),
AxisVector{3, 2, 1, 0}},
{{parameter(f32, {2, 4}), constant<int>(i32, {}, {0}), constant<int>(i32, {1}, {1})},
set_transpose_for,
constant<int>(i32, {1}, {0}),
AxisVector{0}},
{{parameter(f32, {2, 4}), constant<int>(i32, {1}, {0}), constant<int>(i32, {1}, {1})},
set_transpose_for,
constant<int>(i32, {1}, {0}),
AxisVector{1, 0}},
{{parameter(f32, {2, 3, 4}), constant<int>(i32, {2, 3}, {0, 1, 0, 1, 0, 1}), constant<int>(i32, {1}, {1})},
set_transpose_for,
constant<int>(i32, {1}, {1}),
AxisVector{3, 1, 2, 0}}};
AxisVector{3, 1, 2, 0}},
{{parameter(f32, {64, 49, 3, 3, 32}), constant<int>(i32, {}, {1}), constant<int>(i32, {}, {0})},
SetTransposeWithOrder(AxisVector{2, 0, 3, 1, 4}),
constant<int>(i32, {}, {2}),
AxisVector{0, 2, 1, 3}}};
INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_0, TSTestFixture, test_forward_gather(tests_arguments_fw[0]));
INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_1, TSTestFixture, test_forward_gather(tests_arguments_fw[1]));
INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_2, TSTestFixture, test_forward_gather(tests_arguments_fw[2]));
INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_3, TSTestFixture, test_forward_gather(tests_arguments_fw[3]));
INSTANTIATE_TEST_SUITE_P(TSCommonGatherForward_4, TSTestFixture, test_forward_gather(tests_arguments_fw[4]));
struct GatherBackwardArguments {
OutputVector inputs_to_main;

View File

@ -52,6 +52,19 @@ OutputVector set_transpose_for(const vector<size_t>& idxs, const OutputVector& o
return result;
}
OutputVector set_transpose_with_order(const vector<size_t>& idxs,
const OutputVector& out_vec,
const vector<size_t>& transpose_order_axes) {
OutputVector result = out_vec;
for (const auto& idx : idxs) {
const auto& out = out_vec[idx];
auto order = make_shared<Constant>(element::i32, Shape{transpose_order_axes.size()}, transpose_order_axes);
auto transpose = make_shared<Transpose>(out, order);
result[idx] = transpose;
}
return result;
}
OutputVector set_gather_for(const vector<size_t>& idxs, const OutputVector& out_vec) {
OutputVector result = out_vec;
for (const auto& idx : idxs) {

View File

@ -21,6 +21,9 @@ std::string to_string(const ov::Shape& shape);
ov::ParameterVector filter_parameters(const ov::OutputVector& out_vec);
ov::OutputVector set_transpose_for(const std::vector<size_t>& idxs, const ov::OutputVector& out_vec);
ov::OutputVector set_transpose_with_order(const std::vector<size_t>& idxs,
const ov::OutputVector& out_vec,
const std::vector<size_t>& transpose_order_axes);
ov::OutputVector set_gather_for(const std::vector<size_t>& idxs, const ov::OutputVector& out_vec);
std::shared_ptr<ov::Node> create_main_node(const ov::OutputVector& inputs, size_t num_ops, const FactoryPtr& creator);