| 25 | using ::testing::ElementsAreArray; |
| 26 | |
| 27 | class DepthToSpaceOpModel : public SingleOpModel { |
| 28 | public: |
| 29 | DepthToSpaceOpModel(const TensorData& tensor_data, int block_size) { |
| 30 | input_ = AddInput(tensor_data); |
| 31 | output_ = AddOutput(tensor_data); |
| 32 | SetBuiltinOp(BuiltinOperator_DEPTH_TO_SPACE, |
| 33 | BuiltinOptions_DepthToSpaceOptions, |
| 34 | CreateDepthToSpaceOptions(builder_, block_size).Union()); |
| 35 | BuildInterpreter({GetShape(input_)}); |
| 36 | } |
| 37 | |
| 38 | template <typename T> |
| 39 | void SetInput(std::initializer_list<T> data) { |
| 40 | PopulateTensor<T>(input_, data); |
| 41 | } |
| 42 | template <typename T> |
| 43 | std::vector<T> GetOutput() { |
| 44 | return ExtractVector<T>(output_); |
| 45 | } |
| 46 | std::vector<int> GetOutputShape() { return GetTensorShape(output_); } |
| 47 | |
| 48 | private: |
| 49 | int input_; |
| 50 | int output_; |
| 51 | }; |
| 52 | |
| 53 | #ifdef GTEST_HAS_DEATH_TEST |
| 54 | TEST(DepthToSpaceOpModel, BadBlockSize) { |