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: 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> 54 void SetInput(std::initializer_list<T> data) { 55 PopulateTensor<T>(input_, data); 56 } 57 void SetSizeSplits(std::initializer_list<int> data) { 58 PopulateTensor(size_splits_, data); 59 } 60 void SetAxis(int axis) { PopulateTensor(axis_, {axis}); } 61 62 template <typename T> 63 std::vector<T> GetOutput(int i) { 64 return ExtractVector<T>(outputs_[i]); 65 } 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> 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 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 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 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 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 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 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 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 206 int main(int argc, char** argv) { 207 ::tflite::LogToStderr(); 208 ::testing::InitGoogleTest(&argc, argv); 209 return RUN_ALL_TESTS(); 210 } 211
