fix TSGatherForward transformation (#19821)
* fix ts_gather_forward * remove unneeded header
This commit is contained in:
parent
9c908f5245
commit
35d0d92ef7
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue