fix for sliced input, fix for reshaping of Loop op (squash commits) (#3261)

This commit is contained in:
Ivan Tikhonov 2020-11-22 13:25:04 +03:00 committed by GitHub
parent 1c35875a60
commit ec8527f21e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 675 additions and 79 deletions

View File

@ -63,10 +63,11 @@ Loop operation description in the IR has regular sections: `input` and `output`.
Loop operation description in the IR also has several special sections: `body`, `port_map` and `back_edges` similar to the ones from the TensorIterator operation but having some important features described below.
1. The body operation getting an input from the main graph should have an entry in the `port_map` section of the Loop operation. These edges connect input ports of the Loop with the body `Parameter`s.
2. The body operation producing tensor to be used in the subsequent iterations (like in RNN models) should have a back edge described in the `back_edges` section of the operation. The back edge connects the respective body `Parameter` and `Result` operations. For such a case the Loop operation node provides input for the first iteration, while corresponding Loop operation output produces the tensor computed during the last iteration.
3. Output tensors produced by a particular body operation across all iterations can be concatenated and returned as a Loop operation output (this is a "scan output" according to the ONNX* Loop operation [specification](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Loop-13)). The corresponding `output` entry in the `port_map` should have `axis` attribute specifying the axis to concatenate. Therefore, outputs from operations corresponding to `output` entries in the `port_map` without `axis` attribute are returned "as is" (without concatenation).
4. There is one body `Parameter` operation not connected through the `port_map`. This is a "current iteration" input. The Loop operation is responsible for providing the appropriate value for each iteration.
5. Connection of nodes inside the Loop body with the main graph should be done through `Parameter` and `Result` body operations. No other ways to connect graphs are allowed.
2. Input tensors to the Loop can be sliced along a specified axis, the Loop can iterates over all sliced parts. The corresponding `input` entry in the `port_map` should have `axis` attribute specifying the axis to slice. Therefore, inputs to the Loop operation corresponding to `input` entries in the `port_map` without `axis` attribute are used "as is" (without slicing).
3. The body operation producing tensor to be used in the subsequent iterations (like in RNN models) should have a back edge described in the `back_edges` section of the operation. The back edge connects the respective body `Parameter` and `Result` operations. For such a case the Loop operation node provides input for the first iteration, while corresponding Loop operation output produces the tensor computed during the last iteration.
4. Output tensors produced by a particular body operation across all iterations can be concatenated and returned as a Loop operation output (this is a "scan output" according to the ONNX* Loop operation [specification](https://github.com/onnx/onnx/blob/master/docs/Changelog.md#Loop-13)). The corresponding `output` entry in the `port_map` should have `axis` attribute specifying the axis to concatenate. Therefore, outputs from operations corresponding to `output` entries in the `port_map` without `axis` attribute are returned "as is" (without concatenation).
5. There is one body `Parameter` operation not connected through the `port_map`. This is a "current iteration" input. The Loop operation is responsible for providing the appropriate value for each iteration.
6. Connection of nodes inside the Loop body with the main graph should be done through `Parameter` and `Result` body operations. No other ways to connect graphs are allowed.
**Loop attributes**:
@ -101,7 +102,8 @@ Loop operation description in the IR also has several special sections: `body`,
* *axis*
* **Description**: *axis* is an axis to concatenate the body `Result` output across all iterations. Can be specified for `output` entry only.
* **Description**: if *axis* is specified for `output` entry, then it is an axis to concatenate the body `Result` output across all iterations.
If *axis* is specified for `input` entry, then it is an axis to iterate through, it triggers the slicing of the input tensor.
* **Range of values**: an integer. Negative value means counting dimension from the end.
* **Type**: `int`
* **Default value**: None

View File

@ -0,0 +1,259 @@
// Copyright (C) 2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include <gtest/gtest.h>
#include <ngraph/function.hpp>
#include <ngraph/opsets/opset5.hpp>
#include <cpp/ie_cnn_network.h>
using namespace ngraph;
TEST(SmartReshapeTests, TensorIteratorStaticParameters) {
std::shared_ptr<ngraph::Function> f(nullptr);
{
// That which we iterate over
auto X = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
auto Y = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
auto M = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
X->set_friendly_name("X");
Y->set_friendly_name("Y");
M->set_friendly_name("M");
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
auto Xi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Yi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M_body = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto body_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
// Body
auto sum = std::make_shared<ngraph::opset5::Add>(Xi, Yi);
auto Zo = std::make_shared<ngraph::opset5::Multiply>(sum, M_body);
auto body = std::make_shared<ngraph::Function>(OutputVector{Zo, body_condition, sum},
ParameterVector{Xi, Yi, M_body});
auto tensor_iterator = std::make_shared<opset5::TensorIterator>();
tensor_iterator->set_function(body);
tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 2);
tensor_iterator->set_sliced_input(Yi, Y, -1, -1, 1, 0, 1);
tensor_iterator->set_merged_input(M_body, M, Zo);
// Output 0 is last Zo
auto out0 = tensor_iterator->get_iter_value(body_condition, -1);
auto out1 = tensor_iterator->get_iter_value(Zo, -1);
// Output 1 is concat of Zos
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = tensor_iterator->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto out3 = tensor_iterator->get_iter_value(sum, -1);
f = std::make_shared<Function>(OutputVector{out0, out1, out2, out3}, ParameterVector{X, Y, M});
}
InferenceEngine::CNNNetwork network(f);
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({1, 1, 1}));
// concat output (seq len = 1, so it means num_iter = 1)
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({1, 1, 1}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({1, 1, 1}));
ASSERT_NO_THROW(network.reshape({{"X", {32, 1, 10}}, {"Y", {32, 10, 1}}, {"M", {32, 1, 10}}}));
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({32, 1, 10}));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({32, 10, 10}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({32, 1, 1}));
}
TEST(SmartReshapeTests, TensorIteratorDynamicParameters) {
std::shared_ptr<ngraph::Function> f(nullptr);
{
// That which we iterate over
auto X = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
auto Y = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
auto M = std::make_shared<opset5::Parameter>(element::f32, Shape{1, 1, 1});
X->set_friendly_name("X");
Y->set_friendly_name("Y");
M->set_friendly_name("M");
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
auto Xi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Yi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M_body = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto body_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
// Body
auto sum = std::make_shared<ngraph::opset5::Add>(Xi, Yi);
auto Zo = std::make_shared<ngraph::opset5::Multiply>(sum, M_body);
auto body = std::make_shared<ngraph::Function>(OutputVector{Zo, body_condition, sum},
ParameterVector{Xi, Yi, M_body});
auto tensor_iterator = std::make_shared<opset5::TensorIterator>();
tensor_iterator->set_function(body);
tensor_iterator->set_sliced_input(Xi, X, 0, 1, 1, -1, 2);
tensor_iterator->set_sliced_input(Yi, Y, -1, -1, 1, 0, 1);
tensor_iterator->set_merged_input(M_body, M, Zo);
// Output 0 is last Zo
auto out0 = tensor_iterator->get_iter_value(body_condition, -1);
auto out1 = tensor_iterator->get_iter_value(Zo, -1);
// Output 1 is concat of Zos
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = tensor_iterator->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto out3 = tensor_iterator->get_iter_value(sum, -1);
f = std::make_shared<Function>(OutputVector{out0, out1, out2, out3}, ParameterVector{X, Y, M});
}
InferenceEngine::CNNNetwork network(f);
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({1, 1, 1}));
// concat output (seq len = 1, so it means num_iter = 1)
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({1, 1, 1}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({1, 1, 1}));
ASSERT_NO_THROW(network.reshape({{"X", {32, 1, 10}}, {"Y", {32, 10, 1}}, {"M", {32, 1, 10}}}));
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({32, 1, 10}));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({32, 10, 10}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({32, 1, 1}));
}
TEST(SmartReshapeTests, LoopStaticParameters) {
std::shared_ptr<ngraph::Function> f(nullptr);
{
// That which we iterate over
auto X = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Y = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
X->set_friendly_name("X");
Y->set_friendly_name("Y");
M->set_friendly_name("M");
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
auto current_iteration = std::make_shared<opset5::Parameter>(element::i64, Shape{});
auto Xi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Yi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M_body = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto body_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
auto trip_count =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, ngraph::Shape{}, 10);
auto exec_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
// Body
auto sum = std::make_shared<ngraph::opset5::Add>(Xi, Yi);
auto Zo = std::make_shared<ngraph::opset5::Multiply>(sum, M_body);
auto body = std::make_shared<ngraph::Function>(OutputVector{Zo, body_condition, sum},
ParameterVector{Xi, current_iteration, Yi, M_body});
auto loop = std::make_shared<opset5::Loop>(trip_count, exec_condition);
loop->set_function(body);
loop->set_special_body_ports(ngraph::opset5::Loop::SpecialBodyPorts{1, 1});
loop->set_sliced_input(Xi, X, 0, 1, 1, -1, 2);
loop->set_sliced_input(Yi, Y, -1, -1, 1, 0, 1);
loop->set_merged_input(M_body, M, Zo);
// Output 0 is last Zo
auto out0 = loop->get_iter_value(body_condition, -1);
auto out1 = loop->get_iter_value(Zo, -1);
// Output 1 is concat of Zos
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = loop->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto out3 = loop->get_iter_value(sum, -1);
f = std::make_shared<Function>(OutputVector{out0, out1, out2, out3}, ParameterVector{X, Y, M});
}
InferenceEngine::CNNNetwork network(f);
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
ASSERT_NO_THROW(network.reshape({{"X", {32, 1, 10}}, {"Y", {32, 10, 1}}, {"M", {32, 1, 10}}}));
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({32, 1, 10}));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({32, 10, 10}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({32, 1, 1}));
}
TEST(SmartReshapeTests, LoopDynamicParameters) {
std::shared_ptr<ngraph::Function> f(nullptr);
{
// That which we iterate over
auto X = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Y = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
X->set_friendly_name("X");
Y->set_friendly_name("Y");
M->set_friendly_name("M");
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
auto current_iteration = std::make_shared<opset5::Parameter>(element::i64, Shape{});
auto Xi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Yi = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M_body = std::make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto body_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
auto trip_count =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, ngraph::Shape{}, 10);
auto exec_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
// Body
auto sum = std::make_shared<ngraph::opset5::Add>(Xi, Yi);
auto Zo = std::make_shared<ngraph::opset5::Multiply>(sum, M_body);
auto body = std::make_shared<ngraph::Function>(OutputVector{Zo, body_condition, sum},
ParameterVector{Xi, current_iteration, Yi, M_body});
auto loop = std::make_shared<opset5::Loop>(trip_count, exec_condition);
loop->set_function(body);
loop->set_special_body_ports(ngraph::opset5::Loop::SpecialBodyPorts{1, 1});
loop->set_sliced_input(Xi, X, 0, 1, 1, -1, 2);
loop->set_sliced_input(Yi, Y, -1, -1, 1, 0, 1);
loop->set_merged_input(M_body, M, Zo);
// Output 0 is last Zo
auto out0 = loop->get_iter_value(body_condition, -1);
auto out1 = loop->get_iter_value(Zo, -1);
// Output 1 is concat of Zos
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = loop->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto out3 = loop->get_iter_value(sum, -1);
f = std::make_shared<Function>(OutputVector{out0, out1, out2, out3}, ParameterVector{X, Y, M});
}
InferenceEngine::CNNNetwork network(f);
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible(PartialShape::dynamic()));
ASSERT_NO_THROW(network.reshape({{"X", {32, 1, 10}}, {"Y", {32, 10, 1}}, {"M", {32, 1, 10}}}));
ASSERT_TRUE(network.getFunction()->get_results()[0]->get_output_partial_shape(0).compatible({}));
ASSERT_TRUE(network.getFunction()->get_results()[1]->get_output_partial_shape(0).compatible({32, 1, 10}));
// concat output
ASSERT_TRUE(network.getFunction()->get_results()[2]->get_output_partial_shape(0).compatible({32, 10, 10}));
ASSERT_TRUE(network.getFunction()->get_results()[3]->get_output_partial_shape(0).compatible({32, 1, 1}));
}

View File

@ -86,6 +86,10 @@ protected:
using RefBlobGenerator = std::function<InferenceEngine::Blob::Ptr (const InferenceEngine::TensorDesc &info)>;
std::map<std::string, RefBlobGenerator> inputGens, outputGens;
void CreateSlicedLoop(size_t batch_size, size_t num_iteration, InferenceEngine::Precision iePrc,
InferenceEngine::SizeVector& ieShape);
void CreateSlicedLoopDynCondition(size_t batch_size, size_t num_iteration, InferenceEngine::Precision iePrc,
InferenceEngine::SizeVector& ieShape, size_t trip_count);
InferenceEngine::Blob::Ptr GenerateInput(const InferenceEngine::InputInfo &info) const override {
auto found = inputGens.find(info.name());
if (found != inputGens.end()) {
@ -100,7 +104,7 @@ protected:
return LayerTestsCommon::GenerateInput(info);
}
std::vector<std::vector<std::uint8_t>> PredefinedRefs() {
std::vector<std::vector<std::uint8_t>> CalculateRefs() override {
if (outputGens.empty())
return LayerTestsCommon::CalculateRefs();

View File

@ -7,6 +7,7 @@
#include <vector>
#include <memory>
#include <functional>
#include <numeric>
#include "ie_core.hpp"
@ -72,20 +73,20 @@ namespace LayerTestsDefinitions {
types_separate.push_back(el.second);
}
// Example:
/* auto X = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});
auto Y = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});
auto M = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});*/
/* auto X = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});
auto Y = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});
auto M = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::Shape{32, 1, 10});*/
auto params = ngraph::builder::makeParams(ngPrc, inputs_separate);
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
const std::vector<ngraph::PartialShape> body_params_shapes(inputs_separate.size(), ngraph::PartialShape::dynamic());
auto current_iteration = std::make_shared<ngraph::op::Parameter>(ngraph::element::i64, ngraph::Shape{1});
auto current_iteration = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, ngraph::Shape{1});
//Example:
/* auto Xi = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());
auto Yi = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());
auto M_body = std::make_shared<ngraph::op::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());*/
/* auto Xi = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());
auto Yi = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());
auto M_body = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::f32, ngraph::PartialShape::dynamic());*/
ngraph::ParameterVector body_params;
for (const auto &pshape : body_params_shapes) {
@ -147,9 +148,9 @@ namespace LayerTestsDefinitions {
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = loop->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto result0 = std::make_shared<ngraph::op::Result>(out0);
auto result1 = std::make_shared<ngraph::op::Result>(out1);
auto result2 = std::make_shared<ngraph::op::Result>(out2);
auto result0 = std::make_shared<ngraph::opset5::Result>(out0);
auto result1 = std::make_shared<ngraph::opset5::Result>(out1);
auto result2 = std::make_shared<ngraph::opset5::Result>(out2);
function = std::make_shared<ngraph::Function>(ngraph::ResultVector{result0, result1, result2}, params, "loop");
}
@ -179,7 +180,7 @@ namespace LayerTestsDefinitions {
if (is_static)
return std::make_shared<ngraph::opset5::Constant>(prc, shape, value);
auto input = std::make_shared<ngraph::op::Parameter>(prc, shape);
auto input = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
params.push_back(input);
return input;
};
@ -202,10 +203,10 @@ namespace LayerTestsDefinitions {
// Full loop Dynamic exit loop
// n_iter = count n_iter = ex_val
//
auto b_indx = std::make_shared<ngraph::op::Parameter>(ngraph::element::i64, ngraph::Shape{});
auto b_data = std::make_shared<ngraph::op::Parameter>(prc, ngShape);
auto b_indx_cast = std::make_shared<ngraph::op::Convert>(b_indx, prc);
auto b_add = std::make_shared<ngraph::op::Add>(b_data, b_indx_cast, ngraph::op::AutoBroadcastSpec::NUMPY);
auto b_indx = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, ngraph::Shape{});
auto b_data = std::make_shared<ngraph::opset5::Parameter>(prc, ngShape);
auto b_indx_cast = std::make_shared<ngraph::opset5::Convert>(b_indx, prc);
auto b_add = std::make_shared<ngraph::opset5::Add>(b_data, b_indx_cast);
std::shared_ptr<ngraph::Node> b_cond;
if (dynamic_exit == -1) {
@ -317,13 +318,13 @@ namespace LayerTestsDefinitions {
const auto shape = ngraph::Shape{ieShape};
const auto scalarShape = ngraph::Shape{};
auto start = std::make_shared<ngraph::op::Parameter>(prc, shape);
auto count = std::make_shared<ngraph::op::Constant>(ngraph::element::i64, scalarShape, 5);
auto icond = std::make_shared<ngraph::op::Constant>(ngraph::element::boolean, scalarShape, true);
auto start = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto count = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, scalarShape, 5);
auto icond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
// Loop body
auto b_data = std::make_shared<ngraph::op::Parameter>(prc, shape);
auto b_cond = std::make_shared<ngraph::op::Parameter>(ngraph::element::boolean, scalarShape);
auto b_data = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_cond = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::boolean, scalarShape);
auto body = std::make_shared<ngraph::Function>(
ngraph::OutputVector {b_cond, b_data}, // | passthrough body, no data changes
@ -361,14 +362,14 @@ namespace LayerTestsDefinitions {
const auto shape = ngraph::Shape{ieShape};
const auto scalarShape = ngraph::Shape{};
auto start = std::make_shared<ngraph::op::Parameter>(prc, shape);
auto count = std::make_shared<ngraph::op::Constant>(ngraph::element::i64, scalarShape, 5);
auto icond = std::make_shared<ngraph::op::Constant>(ngraph::element::boolean, scalarShape, true);
auto start = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto count = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, scalarShape, 5);
auto icond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
// Loop body
auto b_data = std::make_shared<ngraph::op::Parameter>(prc, shape);
auto b_cond = std::make_shared<ngraph::op::Constant>(ngraph::element::boolean, scalarShape, true);
auto b_iter = std::make_shared<ngraph::op::Parameter>(ngraph::element::i64, scalarShape);
auto b_data = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_cond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
auto b_iter = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, scalarShape);
auto body = std::make_shared<ngraph::Function>(
ngraph::OutputVector {b_cond, b_data},
@ -394,4 +395,182 @@ namespace LayerTestsDefinitions {
Run();
}
void TrivialLoopTest::CreateSlicedLoop(size_t batch_size, size_t num_iteration, InferenceEngine::Precision iePrc,
InferenceEngine::SizeVector& ieShape) {
const auto prc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(iePrc);
const auto scalarShape = ngraph::Shape{};
auto shape = ngraph::Shape{ieShape};
auto to_slice_shape = ngraph::Shape{ieShape};
to_slice_shape[0] = batch_size;
auto to_slice = std::make_shared<ngraph::opset5::Parameter>(prc, to_slice_shape);
auto start = std::make_shared<ngraph::opset5::Constant>(prc, shape, 0);
auto count = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, scalarShape, num_iteration);
auto icond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
// Loop body
auto b_data = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_recu = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_add = std::make_shared<ngraph::opset5::Add>(b_data, b_recu);
auto b_cond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
auto body = std::make_shared<ngraph::Function>(
ngraph::OutputVector {b_cond, b_add},
ngraph::ParameterVector {b_data, b_recu});
auto loop = std::make_shared<ngraph::opset5::Loop>(count, icond);
loop->set_function(body);
loop->set_special_body_ports({-1, 0});
loop->set_sliced_input(b_data, to_slice, 0, 1, 1, -1, 0);
loop->set_merged_input(b_recu, start, b_add);
loop->get_iter_value(b_add, -1);
function = std::make_shared<ngraph::Function>(
ngraph::OutputVector {loop},
ngraph::ParameterVector {to_slice});
}
void TrivialLoopTest::CreateSlicedLoopDynCondition(size_t batch_size, size_t num_iteration, InferenceEngine::Precision iePrc,
InferenceEngine::SizeVector& ieShape, size_t trip_count) {
auto shape = ngraph::Shape{ieShape};
auto to_slice_shape = ngraph::Shape{ieShape};
to_slice_shape[0] = batch_size;
const auto prc = FuncTestUtils::PrecisionUtils::convertIE2nGraphPrc(iePrc);
const auto scalarShape = ngraph::Shape{};
auto to_slice = std::make_shared<ngraph::opset5::Parameter>(prc, to_slice_shape);
auto start = std::make_shared<ngraph::opset5::Constant>(prc, shape, 0);
auto exit_on = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, scalarShape, num_iteration);
auto count = std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, scalarShape, trip_count);
auto icond = std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, scalarShape, true);
// Loop body
auto b_data = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_recu = std::make_shared<ngraph::opset5::Parameter>(prc, shape);
auto b_add = std::make_shared<ngraph::opset5::Add>(b_data, b_recu);
auto b_iter = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, scalarShape);
auto b_exit_on = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, scalarShape);
auto b_cond = std::make_shared<ngraph::opset5::Less>(b_iter, b_exit_on);
auto body = std::make_shared<ngraph::Function>(
ngraph::OutputVector {b_cond, b_add},
ngraph::ParameterVector {b_data, b_recu, b_iter, b_exit_on});
auto loop = std::make_shared<ngraph::opset5::Loop>(count, icond);
loop->set_function(body);
loop->set_special_body_ports({2, 0});
loop->set_sliced_input(b_data, to_slice, 0, 1, 1, -1, 0);
loop->set_invariant_input(b_exit_on, exit_on);
loop->set_merged_input(b_recu, start, b_add);
loop->get_iter_value(b_add, -1);
function = std::make_shared<ngraph::Function>(
ngraph::OutputVector {loop},
ngraph::ParameterVector {to_slice});
}
TEST_P(TrivialLoopTest, AutoSlicingInput_CheckPredefinedValues) {
SKIP_IF_CURRENT_TEST_IS_DISABLED()
InferenceEngine::Precision iePrc;
InferenceEngine::SizeVector ieShape;
std::tie(iePrc, ieShape, targetDevice) = GetParam();
const size_t batch_size = 5;
const size_t num_iteration = 3;
ieShape[0] = 1;
auto ieShape_to_slice = ieShape;
ieShape_to_slice[0] = batch_size;
CreateSlicedLoop(batch_size, num_iteration, iePrc, ieShape);
Run();
// Precalculated ref blobs
auto blob = make_blob_with_precision({iePrc, ieShape_to_slice, InferenceEngine::TensorDesc::getLayoutByDims(ieShape_to_slice)});
blob->allocate();
std::vector<float> seq_raw_data(batch_size);
std::iota(seq_raw_data.begin(), seq_raw_data.end(), 1);
CommonTestUtils::fill_data_with_broadcast(blob, 0, seq_raw_data);
auto blob_ref = make_blob_with_precision({iePrc, ieShape, InferenceEngine::TensorDesc::getLayoutByDims(ieShape)});
blob_ref->allocate();
CommonTestUtils::fill_data_with_broadcast(blob_ref, 0, { num_iteration * (num_iteration + 1) / 2});
inputGens[""] = [&] (InferenceEngine::TensorDesc tdesc) { return blob; };
outputGens[""] = [&] (InferenceEngine::TensorDesc tdesc) { return blob_ref; };
}
TEST_P(TrivialLoopTest, AutoSlicingInputWithDynCondition_CheckPredefinedValues) {
SKIP_IF_CURRENT_TEST_IS_DISABLED()
InferenceEngine::Precision iePrc;
InferenceEngine::SizeVector ieShape;
std::tie(iePrc, ieShape, targetDevice) = GetParam();
// auto slicing size : 5
// trip count limit : 4
// dyn exit after iter : 3
// ---------------------
// should exit after 4 iterations
const size_t batch_size = 5;
const size_t trip_count = 5;
const size_t num_iteration = 3;
ieShape[0] = 1;
auto ieShape_to_slice = ieShape;
ieShape_to_slice[0] = batch_size;
CreateSlicedLoopDynCondition(batch_size, num_iteration, iePrc, ieShape, trip_count);
// Precalculated ref blobs
auto blob = make_blob_with_precision({iePrc, ieShape_to_slice, InferenceEngine::TensorDesc::getLayoutByDims(ieShape_to_slice)});
blob->allocate();
std::vector<float> seq_raw_data(batch_size);
std::iota(seq_raw_data.begin(), seq_raw_data.end(), 1);
CommonTestUtils::fill_data_with_broadcast(blob, 0, seq_raw_data);
auto blob_ref = make_blob_with_precision({iePrc, ieShape, InferenceEngine::TensorDesc::getLayoutByDims(ieShape)});
blob_ref->allocate();
const size_t real_iter = num_iteration + 1;
CommonTestUtils::fill_data_with_broadcast(blob_ref, 0, { real_iter * (real_iter + 1) / 2});
inputGens[""] = [&] (InferenceEngine::TensorDesc tdesc) { return blob; };
outputGens[""] = [&] (InferenceEngine::TensorDesc tdesc) { return blob_ref; };
Run();
}
TEST_P(TrivialLoopTest, AutoSlicingInput_CheckReference) {
SKIP_IF_CURRENT_TEST_IS_DISABLED()
InferenceEngine::Precision iePrc;
InferenceEngine::SizeVector ieShape;
std::tie(iePrc, ieShape, targetDevice) = GetParam();
const size_t batch_size = 5;
const size_t num_iteration = 3;
ieShape[0] = 1;
auto ieShape_to_slice = ieShape;
ieShape_to_slice[0] = batch_size;
CreateSlicedLoop(batch_size, num_iteration, iePrc, ieShape);
Run();
}
TEST_P(TrivialLoopTest, AutoSlicingInputWithDynCondition_CheckReference) {
SKIP_IF_CURRENT_TEST_IS_DISABLED()
InferenceEngine::Precision iePrc;
InferenceEngine::SizeVector ieShape;
std::tie(iePrc, ieShape, targetDevice) = GetParam();
// auto slicing size : 5
// trip count limit : 4
// dyn exit after iter : 3
// ---------------------
// should exit after 4 iterations
const size_t batch_size = 5;
const size_t trip_count = 5;
const size_t num_iteration = 3;
ieShape[0] = 1;
auto ieShape_to_slice = ieShape;
ieShape_to_slice[0] = batch_size;
CreateSlicedLoopDynCondition(batch_size, num_iteration, iePrc, ieShape, trip_count);
Run();
}
} // namespace LayerTestsDefinitions

View File

@ -65,19 +65,6 @@ namespace ngraph
Loop(const Output<Node>& trip_count, const Output<Node>& execution_condition);
int64_t get_num_iterations() const { return m_num_iterations; }
void set_sliced_input(const std::shared_ptr<Parameter>& parameter,
const Output<Node>& value,
int64_t start,
int64_t stride,
int64_t part_size,
int64_t end,
int64_t axis) override
{
NGRAPH_CHECK(false,
"Incorrect type of input. Implicit slicing is not supported in "
"Loop operation.");
}
Output<Node> get_concatenated_slices(const Output<Node>& value,
int64_t start,
int64_t stride,

View File

@ -17,6 +17,7 @@
#include "runtime/reference/loop.hpp"
#include "runtime/reference/concat.hpp"
#include "runtime/reference/function.hpp"
#include "runtime/reference/split.hpp"
namespace ngraph
{
@ -123,6 +124,47 @@ namespace ngraph
concat_outputs.push_back(concat_desc);
}
}
// Slicing
std::vector<std::shared_ptr<opset5::TensorIterator::SliceInputDescription>>
slice_inputs;
std::vector<HostTensorVector> sliced_values;
int slice_in_idx = 0;
for (const auto& desc : input_descs)
{
if (const auto& slice_desc = std::dynamic_pointer_cast<
opset5::TensorIterator::SliceInputDescription>(desc))
{
const auto el_size =
args[slice_desc->m_input_index]->get_element_type().size();
slice_inputs.push_back(slice_desc);
auto shape = args[slice_desc->m_input_index]->get_shape();
uint64_t num_iterations = shape.at(slice_desc->m_axis);
shape.at(slice_desc->m_axis) = 1;
sliced_values.emplace_back(HostTensorVector());
for (int i = 0; i < num_iterations; ++i)
{
sliced_values.back().emplace_back(std::make_shared<HostTensor>(
args[slice_desc->m_input_index]->get_element_type(), shape));
}
std::vector<char*> pointers_to_data(num_iterations);
for (size_t j = 0; j < pointers_to_data.size(); ++j)
{
pointers_to_data[slice_desc->m_stride > 0
? j
: (pointers_to_data.size() - j - 1)] =
sliced_values[slice_in_idx][j]->get_data_ptr<char>();
}
reference::split(args[slice_desc->m_input_index]->get_data_ptr<char>(),
args[slice_desc->m_input_index]->get_shape(),
el_size,
slice_desc->m_axis,
num_iterations,
pointers_to_data.data());
slice_in_idx++;
}
}
// Allocate vectors for store output values
std::vector<HostTensorVector> values_to_concat(concat_outputs.size());
HostTensorVector body_outputs;
@ -131,6 +173,13 @@ namespace ngraph
trip_count = trip_count >= 0 ? trip_count : std::numeric_limits<int64_t>::max();
for (int64_t cur_iter = 0; cur_iter < trip_count; ++cur_iter)
{
// Copy new values for sliced inputs
for (size_t i = 0; i < slice_inputs.size(); ++i)
{
inputs_to_body[slice_inputs[i]->m_body_parameter_index] =
sliced_values[i][cur_iter];
}
// Evaluate body
body_outputs.clear();
reference::function(func, inputs_to_body, body_outputs);

View File

@ -172,8 +172,6 @@ void op::v5::Loop::validate_and_infer_types()
get_output_size() == m_output_descriptions.size(),
"Number of outputs must be the same as number of output descriptions");
std::vector<std::shared_ptr<Node>> ends;
// Input
uint64_t index_it = 2;
for (const auto& input_description : m_input_descriptions)
@ -182,11 +180,29 @@ void op::v5::Loop::validate_and_infer_types()
NODE_VALIDATION_CHECK(this, index == index_it, "Input_index not in order");
index_it++;
if (auto merged_input_description = as_type_ptr<MergedInputDescription>(input_description))
if (auto slice_input_description = as_type_ptr<SliceInputDescription>(input_description))
{
auto body_parameter =
m_body->get_parameters().at(slice_input_description->m_body_parameter_index);
auto input_partial_shape = inputs().at(index).get_source_output().get_partial_shape();
if (input_partial_shape.is_static())
{
// infer type for m_body_parameter
Shape out_shape{input_partial_shape.to_shape()};
out_shape[slice_input_description->m_axis] = slice_input_description->m_part_size;
body_parameter->set_partial_shape(out_shape);
}
else
{
body_parameter->set_partial_shape(
PartialShape::dynamic(input_partial_shape.rank()));
}
}
else if (auto merged_input_description =
as_type_ptr<MergedInputDescription>(input_description))
{
auto body_value =
m_body->get_results().at(merged_input_description->m_body_value_index);
ends.push_back(body_value);
const auto& body_value_partial_shape = body_value->get_input_partial_shape(0);
auto body_parameter =
@ -194,22 +210,8 @@ void op::v5::Loop::validate_and_infer_types()
auto body_param_partial_shape = body_parameter->get_partial_shape();
auto input_partial_shape = input(index).get_partial_shape();
NODE_VALIDATION_CHECK(this,
body_value_partial_shape.compatible(body_param_partial_shape),
"Iterator successive value is not compatible with body param");
NODE_VALIDATION_CHECK(this,
input_partial_shape.compatible(body_param_partial_shape),
"Iterator initial value is not compatible with body param");
if (input_partial_shape.is_static())
{
auto input_shape = input_partial_shape.to_shape();
// infer type for body_parameter
if (body_param_partial_shape.is_dynamic())
{
body_parameter->set_partial_shape(input_shape);
}
}
body_parameter->set_partial_shape(input_partial_shape);
}
else if (auto invariant_input_description =
as_type_ptr<TensorIterator::InvariantInputDescription>(input_description))
@ -223,15 +225,7 @@ void op::v5::Loop::validate_and_infer_types()
input_partial_shape.compatible(body_param_partial_shape),
"Iterator initial value is not compatible with body param");
if (input_partial_shape.is_static())
{
auto input_shape = input_partial_shape.to_shape();
// infer type for m_body_parameter
if (body_param_partial_shape.is_dynamic())
{
body_parameter->set_partial_shape(input_shape);
}
}
body_parameter->set_partial_shape(input_partial_shape);
}
}
@ -282,6 +276,12 @@ void op::v5::Loop::validate_and_infer_types()
set_output_type(index, body_value.get_element_type(), out_shape);
}
}
else
{
set_output_type(index,
body_value.get_element_type(),
PartialShape::dynamic(body_value.get_partial_shape().rank()));
}
}
else if (auto body_output_description =
as_type_ptr<TensorIterator::BodyOutputDescription>(output_description))
@ -334,6 +334,17 @@ std::shared_ptr<Node> op::v5::Loop::clone_with_new_inputs(const OutputVector& ne
new_args[input_index].get_element_type();
new_shapes[input_description->m_body_parameter_index] =
new_args[input_index].get_partial_shape();
if (new_shapes[input_description->m_body_parameter_index].is_static())
{
if (auto slice_in = ::ngraph::as_type_ptr<
ngraph::op::v0::TensorIterator::SliceInputDescription>(
input_description))
{
new_shapes[slice_in->m_body_parameter_index][slice_in->m_axis] =
slice_in->m_part_size;
}
}
}
}
}

View File

@ -152,12 +152,6 @@ void op::v0::TensorIterator::validate_and_infer_types()
auto body_param_partial_shape = body_parameter->get_partial_shape();
auto input_partial_shape = inputs().at(index).get_source_output().get_partial_shape();
NODE_VALIDATION_CHECK(this,
body_value_partial_shape.compatible(body_param_partial_shape),
"Iterator successive value is not compatible with body param");
NODE_VALIDATION_CHECK(this,
input_partial_shape.compatible(body_param_partial_shape),
"Iterator initial value is not compatible with body param");
body_parameter->set_partial_shape(input_partial_shape);
}
else if (auto invariant_input_description =

View File

@ -751,3 +751,114 @@ TEST(type_prop, loop_operation_for_mode_10_iter_static_shapes_special_body_ports
EXPECT_EQ(loop->get_output_shape(1), out1_shape);
EXPECT_EQ(loop->get_output_shape(2), out2_shape);
}
// SliceInputs testing
// trip_count = 10
// execution_condition = true
// body_condition = true
// all shapes are static, 10 iterations will be executed
TEST(type_prop, loop_operation_10_iter_static_shapes_sliced_inputs)
{
// That which we iterate over
auto X = make_shared<opset5::Parameter>(element::f32, Shape{32, 1, 10});
auto Y = make_shared<opset5::Parameter>(element::f32, Shape{32, 10, 1});
auto M = make_shared<opset5::Parameter>(element::f32, Shape{32, 1, 10});
// Set up the cell body, a function from (Xi, Yi) -> (Zo)
// Body parameters
auto current_iteration = make_shared<opset5::Parameter>(element::i64, Shape{});
auto Xi = make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto Yi = make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto M_body = make_shared<opset5::Parameter>(element::f32, PartialShape::dynamic());
auto body_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
auto trip_count =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::i64, ngraph::Shape{}, 10);
auto exec_condition =
std::make_shared<ngraph::opset5::Constant>(ngraph::element::boolean, ngraph::Shape{}, true);
// Body
auto sum = make_shared<ngraph::opset5::Add>(Xi, Yi);
auto Zo = make_shared<ngraph::opset5::Multiply>(sum, M_body);
auto body = make_shared<ngraph::Function>(OutputVector{Zo, body_condition, sum},
ParameterVector{Xi, current_iteration, Yi, M_body});
auto loop = make_shared<opset5::Loop>(trip_count, exec_condition);
loop->set_function(body);
loop->set_special_body_ports(ngraph::opset5::Loop::SpecialBodyPorts{1, 1});
loop->set_sliced_input(Xi, X, 0, 1, 1, -1, 2);
loop->set_sliced_input(Yi, Y, -1, -1, 1, 0, 1);
loop->set_merged_input(M_body, M, Zo);
// check input descriptors
for (auto& desc : loop->get_input_descriptions())
{
auto type_info = desc->get_type_info();
if (std::strcmp(type_info.name, "InvariantInputDescription") == 0)
{
auto input_desc =
as_type_ptr<ngraph::opset5::TensorIterator::InvariantInputDescription>(desc);
EXPECT_NE(input_desc, nullptr);
}
else if (std::strcmp(type_info.name, "SliceInputDescription") == 0)
{
auto input_desc =
as_type_ptr<ngraph::opset5::TensorIterator::SliceInputDescription>(desc);
EXPECT_NE(input_desc, nullptr);
}
else if (std::strcmp(type_info.name, "MergedInputDescription") == 0)
{
auto input_desc =
as_type_ptr<ngraph::opset5::TensorIterator::MergedInputDescription>(desc);
EXPECT_NE(input_desc, nullptr);
}
}
// Output 0 is last Zo
auto out0 = loop->get_iter_value(body_condition, -1);
auto out1 = loop->get_iter_value(Zo, -1);
// Output 1 is concat of Zos
// start=0, stride=1, part_size=1, end=-1, axis=1
auto out2 = loop->get_concatenated_slices(Zo, 0, 1, 1, -1, 1);
auto out3 = loop->get_iter_value(sum, -1);
// check output descriptors
for (auto& desc : loop->get_output_descriptions())
{
auto type_info = desc->get_type_info();
if (std::strcmp(type_info.name, "ConcatOutputDescription") == 0)
{
auto output_desc =
as_type_ptr<ngraph::opset5::TensorIterator::ConcatOutputDescription>(desc);
EXPECT_NE(output_desc, nullptr);
}
else if (std::strcmp(type_info.name, "BodyOutputDescription") == 0)
{
auto output_desc =
as_type_ptr<ngraph::opset5::TensorIterator::BodyOutputDescription>(desc);
EXPECT_NE(output_desc, nullptr);
}
}
auto result0 = make_shared<opset5::Result>(out0);
auto result1 = make_shared<opset5::Result>(out1);
auto result2 = make_shared<opset5::Result>(out2);
auto result3 = make_shared<opset5::Result>(out3);
Shape out0_shape{};
Shape out1_shape{32, 1, 10};
Shape out2_shape{32, 10, 10};
Shape out3_shape{32, 1, 1};
auto results = ResultVector{result0, result1, result2, result3};
auto f = make_shared<Function>(results, ParameterVector{X, Y, M});
EXPECT_EQ(result0->get_output_shape(0), out0_shape);
EXPECT_EQ(result1->get_output_shape(0), out1_shape);
EXPECT_EQ(result2->get_output_shape(0), out2_shape);
EXPECT_EQ(result3->get_output_shape(0), out3_shape);
EXPECT_EQ(loop->get_output_shape(0), out0_shape);
EXPECT_EQ(loop->get_output_shape(1), out1_shape);
EXPECT_EQ(loop->get_output_shape(2), out2_shape);
EXPECT_EQ(loop->get_output_shape(3), out3_shape);
}