Internal change on flatbuffer implementation.

PiperOrigin-RevId: 327150007
Change-Id: If35620d1b9556269f3d4a26443ecbd1130de30df
This commit is contained in:
Hyeonjong Ryu 2020-08-17 19:45:48 -07:00 committed by TensorFlower Gardener
parent 70b61c44ae
commit 581ef323c7
2 changed files with 10 additions and 9 deletions

View File

@ -633,10 +633,14 @@ TfLiteStatus ParseOpDataTfLite(const Operator* op, BuiltinOperator op_type,
TF_LITE_ENSURE(error_reporter, params != nullptr);
if (const auto* schema_params = op->builtin_options_as_SqueezeOptions()) {
const auto* squeeze_dims = schema_params->squeeze_dims();
TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray(
sizeof(params->squeeze_dims), squeeze_dims, params->squeeze_dims,
error_reporter, "squeeze"));
params->num_squeeze_dims = squeeze_dims->size();
if (squeeze_dims != nullptr) {
TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray(
sizeof(params->squeeze_dims), squeeze_dims, params->squeeze_dims,
error_reporter, "squeeze"));
params->num_squeeze_dims = squeeze_dims->size();
} else {
params->num_squeeze_dims = 0;
}
}
*builtin_data = params.release();
return kTfLiteOk;

View File

@ -82,15 +82,12 @@ class FlatbufferConversionsTest : public ::testing::Test {
flatbuffers::FlatBufferBuilder builder_;
};
TEST_F(FlatbufferConversionsTest, ParseBadSqueeze) {
TEST_F(FlatbufferConversionsTest, ParseSqueezeAll) {
const Operator* op = BuildTestOperator(
BuiltinOptions_SqueezeOptions, CreateSqueezeOptions(builder_).Union());
void* output_data = nullptr;
EXPECT_NE(kTfLiteOk, ParseOpData(op, BuiltinOperator_SQUEEZE, &mock_reporter_,
EXPECT_EQ(kTfLiteOk, ParseOpData(op, BuiltinOperator_SQUEEZE, &mock_reporter_,
&mock_allocator_, &output_data));
EXPECT_THAT(mock_reporter_.GetAsString(),
::testing::ContainsRegex(
"Input array not provided for operation 'squeeze'"));
}
TEST_F(FlatbufferConversionsTest, ParseDynamicReshape) {