CherryPick2.4:Add depth_to_space TFLite op

This commit is contained in:
Mihai Maruseac 2021-04-27 17:47:36 -07:00 committed by Geeta Chavan
parent 29cccf5062
commit c5736e98e3
2 changed files with 6 additions and 0 deletions

View File

@ -61,6 +61,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type); TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
const int block_size = params->block_size; const int block_size = params->block_size;
TF_LITE_ENSURE(context, block_size > 0);
const int input_height = input->dims->data[1]; const int input_height = input->dims->data[1];
const int input_width = input->dims->data[2]; const int input_width = input->dims->data[2];
const int input_channels = input->dims->data[3]; const int input_channels = input->dims->data[3];

View File

@ -60,6 +60,11 @@ TEST(DepthToSpaceOpModel, BadBlockSize) {
EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4), EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4),
"Cannot allocate tensors"); "Cannot allocate tensors");
} }
TEST(DepthToSpaceOpModel, NoBlockSize) {
EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
"Cannot allocate tensors");
}
#endif #endif
TEST(DepthToSpaceOpModel, Float32) { TEST(DepthToSpaceOpModel, Float32) {