1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include <initializer_list>
16 #include <gtest/gtest.h>
17 #include "tensorflow/lite/interpreter.h"
18 #include "tensorflow/lite/kernels/register.h"
19 #include "tensorflow/lite/kernels/test_util.h"
20 #include "tensorflow/lite/model.h"
21
22 namespace tflite {
23 namespace {
24
25 using ::testing::ElementsAreArray;
26
27 constexpr int kAxisIsATensor = -1000;
28
29 class SplitVOpModel : public SingleOpModel {
30 public:
SplitVOpModel(const TensorData & input,const TensorData & size_splits,int num_splits,int axis)31 SplitVOpModel(const TensorData& input, const TensorData& size_splits,
32 int num_splits, int axis) {
33 input_ = AddInput(input);
34 size_splits_ = AddInput(size_splits);
35 if (axis == kAxisIsATensor) {
36 axis_ = AddInput({TensorType_INT32, {1}});
37 } else {
38 axis_ = AddConstInput(TensorType_INT32, {axis}, {1});
39 }
40 for (int i = 0; i < num_splits; ++i) {
41 outputs_.push_back(AddOutput(input.type));
42 }
43 SetBuiltinOp(BuiltinOperator_SPLIT_V, BuiltinOptions_SplitVOptions,
44 CreateSplitVOptions(builder_, num_splits).Union());
45 if (axis == kAxisIsATensor) {
46 BuildInterpreter(
47 {GetShape(input_), GetShape(size_splits_), GetShape(axis_)});
48 } else {
49 BuildInterpreter({GetShape(input_), GetShape(size_splits_), {}});
50 }
51 }
52
53 template <typename T>
SetInput(std::initializer_list<T> data)54 void SetInput(std::initializer_list<T> data) {
55 PopulateTensor<T>(input_, data);
56 }
SetSizeSplits(std::initializer_list<int> data)57 void SetSizeSplits(std::initializer_list<int> data) {
58 PopulateTensor(size_splits_, data);
59 }
SetAxis(int axis)60 void SetAxis(int axis) { PopulateTensor(axis_, {axis}); }
61
62 template <typename T>
GetOutput(int i)63 std::vector<T> GetOutput(int i) {
64 return ExtractVector<T>(outputs_[i]);
65 }
GetOutputShape(int i)66 std::vector<int> GetOutputShape(int i) { return GetTensorShape(outputs_[i]); }
67
68 private:
69 int input_;
70 int size_splits_;
71 int axis_;
72 std::vector<int> outputs_;
73 };
74
75 template <typename T, TensorType T1>
Check(int axis,std::initializer_list<int> input_shape,std::initializer_list<int> size_splits_shape,std::vector<std::initializer_list<int>> output_shapes,const std::initializer_list<T> & input_data,const std::initializer_list<int> & size_splits_data,const std::vector<std::initializer_list<T>> & output_data)76 void Check(int axis, std::initializer_list<int> input_shape,
77 std::initializer_list<int> size_splits_shape,
78 std::vector<std::initializer_list<int>> output_shapes,
79 const std::initializer_list<T>& input_data,
80 const std::initializer_list<int>& size_splits_data,
81 const std::vector<std::initializer_list<T>>& output_data) {
82 int num_splits = size_splits_data.size();
83 SplitVOpModel m({T1, input_shape}, {TensorType_INT32, size_splits_shape},
84 num_splits, kAxisIsATensor);
85 m.SetInput<T>(input_data);
86 m.SetSizeSplits(size_splits_data);
87 m.SetAxis(axis);
88 m.Invoke();
89 for (int i = 0; i < num_splits; ++i) {
90 EXPECT_THAT(m.GetOutput<T>(i), ElementsAreArray(output_data[i]));
91 EXPECT_THAT(m.GetOutputShape(i), ElementsAreArray(output_shapes[i]));
92 }
93
94 SplitVOpModel const_m({T1, input_shape},
95 {TensorType_INT32, size_splits_shape}, num_splits,
96 axis);
97 const_m.SetInput<T>(input_data);
98 const_m.SetSizeSplits(size_splits_data);
99 const_m.Invoke();
100 for (int i = 0; i < num_splits; ++i) {
101 EXPECT_THAT(const_m.GetOutput<T>(i), ElementsAreArray(output_data[i]));
102 EXPECT_THAT(const_m.GetOutputShape(i), ElementsAreArray(output_shapes[i]));
103 }
104 }
105
TEST(SplitVOpTest,TwoDimensional)106 TEST(SplitVOpTest, TwoDimensional) {
107 // Input shape: {4, 3}
108 // size_splits: {1, 1, 2}
109 // axis: 0
110 // We should have 3 outpus with shapes respectively:
111 // output 1 : {1, 3}
112 // output 2 : {1, 3}
113 // output 3 : {2, 3}
114 Check<float, TensorType_FLOAT32>(
115 /*axis=*/0, {4, 3}, {3}, {{1, 3}, {1, 3}, {2, 3}},
116 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, {1, 1, 2},
117 {{1, 2, 3}, {4, 5, 6}, {7, 8, 9, 10, 11, 12}});
118 }
119
TEST(SplitVOpTest,FourDimensional)120 TEST(SplitVOpTest, FourDimensional) {
121 Check<float, TensorType_FLOAT32>(
122 /*axis=*/0, {2, 2, 2, 2}, {2}, {{1, 2, 2, 2}, {1, 2, 2, 2}},
123 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {1, 1},
124 {
125 {1, 2, 3, 4, 5, 6, 7, 8},
126 {9, 10, 11, 12, 13, 14, 15, 16},
127 });
128 Check<float, TensorType_FLOAT32>(
129 /*axis=*/1, {2, 2, 2, 2}, {2}, {{2, 1, 2, 2}, {2, 1, 2, 2}},
130 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {1, -1},
131 {
132 {1, 2, 3, 4, 9, 10, 11, 12},
133 {5, 6, 7, 8, 13, 14, 15, 16},
134 });
135 Check<float, TensorType_FLOAT32>(
136 /*axis=*/2, {2, 2, 2, 2}, {2}, {{2, 2, 1, 2}, {2, 2, 1, 2}},
137 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {1, 1},
138 {
139 {1, 2, 5, 6, 9, 10, 13, 14},
140 {3, 4, 7, 8, 11, 12, 15, 16},
141 });
142 Check<float, TensorType_FLOAT32>(
143 /*axis=*/3, {2, 2, 2, 2}, {2}, {{2, 2, 2, 1}, {2, 2, 2, 1}},
144 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {1, 1},
145 {
146 {1, 3, 5, 7, 9, 11, 13, 15},
147 {2, 4, 6, 8, 10, 12, 14, 16},
148 });
149 }
150
TEST(SplitVOpTest,OneDimensional)151 TEST(SplitVOpTest, OneDimensional) {
152 Check<float, TensorType_FLOAT32>(
153 /*axis=*/0, {8}, {8}, {{1}, {1}, {1}, {1}, {1}, {1}, {1}, {1}},
154 {1, 2, 3, 4, 5, 6, 7, 8}, {1, 1, 1, 1, 1, 1, 1, 1},
155 {{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}});
156 }
157
TEST(SplitVOpTest,OneDimensional2)158 TEST(SplitVOpTest, OneDimensional2) {
159 Check<float, TensorType_FLOAT32>(
160 /*axis=*/0, {8}, {8}, {{1}, {1}, {1}, {1}, {1}, {1}, {2}, {0}},
161 {1, 2, 3, 4, 5, 6, 7, 8}, {1, 1, 1, 1, 1, 1, 2, -1},
162 {{1}, {2}, {3}, {4}, {5}, {6}, {7, 8}, {}});
163 }
164
TEST(SplitVOpTest,NegativeAxis)165 TEST(SplitVOpTest, NegativeAxis) {
166 Check<float, TensorType_FLOAT32>(
167 /*axis=*/-4, {2, 2, 2, 2}, {2}, {{1, 2, 2, 2}, {1, 2, 2, 2}},
168 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {1, 1},
169 {
170 {1, 2, 3, 4, 5, 6, 7, 8},
171 {9, 10, 11, 12, 13, 14, 15, 16},
172 });
173 }
174
TEST(SplitVOpTest,TwoDimensionalUint8)175 TEST(SplitVOpTest, TwoDimensionalUint8) {
176 // Input shape: {4, 3}
177 // size_splits: {1, 1, 2}
178 // axis: 0
179 // We should have 3 outpus with shapes respectively:
180 // output 1 : {1, 3}
181 // output 2 : {1, 3}
182 // output 3 : {2, 3}
183 Check<uint8_t, TensorType_UINT8>(
184 /*axis=*/0, {4, 3}, {3}, {{1, 3}, {1, 3}, {2, 3}},
185 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, {1, 1, 2},
186 {{1, 2, 3}, {4, 5, 6}, {7, 8, 9, 10, 11, 12}});
187 }
188
TEST(SplitVOpTest,TwoDimensionalInt16)189 TEST(SplitVOpTest, TwoDimensionalInt16) {
190 // Input shape: {4, 3}
191 // size_splits: {1, 1, 2}
192 // axis: 0
193 // We should have 3 outpus with shapes respectively:
194 // output 1 : {1, 3}
195 // output 2 : {1, 3}
196 // output 3 : {2, 3}
197 Check<int16_t, TensorType_INT16>(
198 /*axis=*/0, {4, 3}, {3}, {{1, 3}, {1, 3}, {2, 3}},
199 {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, {1, 1, 2},
200 {{1, 2, 3}, {4, 5, 6}, {7, 8, 9, 10, 11, 12}});
201 }
202
203 } // namespace
204 } // namespace tflite
205
main(int argc,char ** argv)206 int main(int argc, char** argv) {
207 ::tflite::LogToStderr();
208 ::testing::InitGoogleTest(&argc, argv);
209 return RUN_ALL_TESTS();
210 }
211