Treat 0-size constants properly as constants in InterpreterBuilder.
PiperOrigin-RevId: 468311951
diff --git a/tensorflow/lite/BUILD b/tensorflow/lite/BUILD
index 2fed190..fa29d6b 100644
--- a/tensorflow/lite/BUILD
+++ b/tensorflow/lite/BUILD
@@ -984,6 +984,7 @@
"testdata/test_model_broken.bin",
"testdata/test_model_redux_precision.bin",
"testdata/while_op_with_forwarding_input.bin",
+ "testdata/zero_size_constant.bin",
],
tags = [
"no_windows", # TODO(b/194459105): the test is flaky.
@@ -994,7 +995,8 @@
":framework",
":interpreter_test_util",
":string",
- "//tensorflow/lite:string_util",
+ ":string_util",
+ "//tensorflow/lite/c:common",
"//tensorflow/lite/core/api",
"//tensorflow/lite/core/api:verifier",
"//tensorflow/lite/kernels:builtin_ops",
diff --git a/tensorflow/lite/interpreter_builder.cc b/tensorflow/lite/interpreter_builder.cc
index 3c70a90..5473103 100644
--- a/tensorflow/lite/interpreter_builder.cc
+++ b/tensorflow/lite/interpreter_builder.cc
@@ -603,11 +603,9 @@
}
if (auto* buffer = (*buffers)[tensor->buffer()]) {
if (auto* array = buffer->data()) {
- if (size_t size = array->size()) {
- *buffer_size = size;
- *buffer_data = reinterpret_cast<const char*>(array->data());
- return kTfLiteOk;
- }
+ *buffer_size = array->size();
+ *buffer_data = reinterpret_cast<const char*>(array->data());
+ return kTfLiteOk;
}
}
return kTfLiteOk;
diff --git a/tensorflow/lite/model_test.cc b/tensorflow/lite/model_test.cc
index 1cd3523..bb0e602 100644
--- a/tensorflow/lite/model_test.cc
+++ b/tensorflow/lite/model_test.cc
@@ -29,6 +29,7 @@
#include <gtest/gtest.h>
#include "flatbuffers/flatbuffers.h" // from @flatbuffers
#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/api/error_reporter.h"
#include "tensorflow/lite/core/api/op_resolver.h"
#include "tensorflow/lite/core/api/verifier.h"
@@ -789,6 +790,31 @@
ASSERT_EQ(interpreter->Invoke(), kTfLiteOk);
}
+TEST(BasicFlatBufferModel, TestHandleZeroSizeConstant) {
+ TestErrorReporter reporter;
+ FileCopyAllocation model_allocation(
+ "tensorflow/lite/testdata/zero_size_constant.bin", &reporter);
+ EXPECT_TRUE(model_allocation.valid());
+ ::flatbuffers::Verifier verifier(
+ reinterpret_cast<const uint8_t*>(model_allocation.base()),
+ model_allocation.bytes());
+ EXPECT_TRUE(VerifyModelBuffer(verifier));
+ const Model* model_fb = ::tflite::GetModel(model_allocation.base());
+
+ auto model = FlatBufferModel::BuildFromModel(model_fb);
+ EXPECT_TRUE(model);
+
+ std::unique_ptr<Interpreter> interpreter;
+ EXPECT_EQ(
+ InterpreterBuilder(*model, TrivialResolver(&dummy_reg))(&interpreter),
+ kTfLiteOk);
+ EXPECT_NE(interpreter, nullptr);
+
+ EXPECT_EQ(interpreter->tensors_size(), 3);
+ // Second tensor should be treated as constant.
+ ASSERT_EQ(interpreter->tensor(1)->allocation_type, kTfLiteMmapRo);
+}
+
// TODO(aselle): Add tests for serialization of builtin op data types.
// These tests will occur with the evaluation tests of individual operators,
// not here.
diff --git a/tensorflow/lite/testdata/zero_size_constant.bin b/tensorflow/lite/testdata/zero_size_constant.bin
new file mode 100644
index 0000000..ac28ffb
--- /dev/null
+++ b/tensorflow/lite/testdata/zero_size_constant.bin
Binary files differ
