[lite]

Set the bytes field by default in case no resize happened while assigning to resource variable.
Also, add unit-test for the resource variable

PiperOrigin-RevId: 350585904
Change-Id: Ib3e74d5aca21abf0cf1dd227228b05744e680caa

tensorflow/lite/experimental/resource/BUILD

@@ -26,3 +26,17 @@ cc_library(
         "//tensorflow/lite/kernels/internal:tensor",
     ],
 )
+
+cc_test(
+    name = "resource_variable_test",
+    srcs = [
+        "resource_variable_test.cc",
+    ],
+    deps = [
+        ":resource",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:common",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/resource/resource_variable.cc

@@ -69,6 +69,8 @@ TfLiteStatus ResourceVariable::AssignFrom(const TfLiteTensor* tensor) {
   tensor_.data.raw = old_raw;
   if (old_bytes != tensor->bytes) {
     TfLiteTensorRealloc(tensor->bytes, &tensor_);
+  } else {
+    tensor_.bytes = old_bytes;
   }
 
   memcpy(tensor_.data.raw, tensor->data.raw, tensor_.bytes);
@@ -82,8 +84,8 @@ void CreateResourceVariableIfNotAvailable(ResourceMap* resources,
   if (resources->count(resource_id) != 0) {
     return;
   }
-  resources->emplace(
-      resource_id, std::unique_ptr<ResourceVariable>(new ResourceVariable()));
+  resources->emplace(resource_id,
+                     std::unique_ptr<ResourceVariable>(new ResourceVariable()));
 }
 
 ResourceVariable* GetResourceVariable(ResourceMap* resources, int resource_id) {

tensorflow/lite/experimental/resource/resource_variable_test.cc

@@ -0,0 +1,173 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/experimental/resource/resource_variable.h"
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace resource {
+// Helper util that initialize 'tensor'.
+void InitTensor(const std::vector<int>& shape, TfLiteAllocationType alloc_type,
+                float default_value, TfLiteTensor* tensor) {
+  memset(tensor, 0, sizeof(TfLiteTensor));
+  int num_elements = 1;
+  for (auto dim : shape) num_elements *= dim;
+  if (shape.empty()) num_elements = 0;
+  float* buf = static_cast<float*>(malloc(sizeof(float) * num_elements));
+  for (int i = 0; i < num_elements; ++i) buf[i] = default_value;
+  const int bytes = num_elements * sizeof(buf[0]);
+  auto* dims = ConvertArrayToTfLiteIntArray(shape.size(), shape.data());
+  TfLiteTensorReset(TfLiteType::kTfLiteFloat32, nullptr, dims, {},
+                    reinterpret_cast<char*>(buf), bytes, alloc_type, nullptr,
+                    false, tensor);
+}
+
+TEST(ResourceTest, NonDynamicTensorAssign) {
+  ResourceVariable var;
+  EXPECT_FALSE(var.IsInitialized());
+
+  TfLiteTensor tensor;
+  std::vector<int> shape = {1};
+  InitTensor(shape, kTfLiteArenaRw, 1.0f, &tensor);
+
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor));
+  EXPECT_TRUE(var.IsInitialized());
+  auto* value = var.GetTensor();
+
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float), value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(1, value->dims->data[0]);
+  EXPECT_EQ(1.0f, value->data.f[0]);
+
+  // Cleanup
+  // For non dynamic tensors we need to delete the buffers manually.
+  free(tensor.data.raw);
+  TfLiteTensorFree(&tensor);
+}
+
+TEST(ResourceTest, DynamicTensorAssign) {
+  ResourceVariable var;
+  EXPECT_FALSE(var.IsInitialized());
+
+  TfLiteTensor tensor;
+  std::vector<int> shape = {1};
+  InitTensor(shape, kTfLiteDynamic, 1.0f, &tensor);
+
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor));
+  EXPECT_TRUE(var.IsInitialized());
+  auto* value = var.GetTensor();
+
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float), value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(1, value->dims->data[0]);
+  EXPECT_EQ(1.0f, value->data.f[0]);
+
+  // Cleanup
+  TfLiteTensorFree(&tensor);
+}
+
+TEST(ResourceTest, AssignSameSizeTensor) {
+  ResourceVariable var;
+  EXPECT_FALSE(var.IsInitialized());
+
+  // We create 2 tensors and make 2 calls for Assign.
+  // The second Assign call should trigger the case of assign with same size.
+  TfLiteTensor tensor_a, tensor_b;
+  std::vector<int> shape_a = {1};
+  std::vector<int> shape_b = {1};
+  InitTensor(shape_a, kTfLiteDynamic, 1.0, &tensor_a);
+  InitTensor(shape_b, kTfLiteDynamic, 4.0, &tensor_b);
+
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_a));
+  EXPECT_TRUE(var.IsInitialized());
+  auto* value = var.GetTensor();
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float), value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(1, value->dims->data[0]);
+  EXPECT_EQ(1.0f, value->data.f[0]);
+
+  // Second AssignFrom but now tensor_b has same size as the variable.
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_b));
+  EXPECT_TRUE(var.IsInitialized());
+  value = var.GetTensor();
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float), value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(1, value->dims->data[0]);
+  EXPECT_EQ(4.0f, value->data.f[0]);
+
+  // Cleanup
+  TfLiteTensorFree(&tensor_a);
+  TfLiteTensorFree(&tensor_b);
+}
+
+TEST(ResourceTest, AssignDifferentSizeTensor) {
+  ResourceVariable var;
+  EXPECT_FALSE(var.IsInitialized());
+
+  // We create 2 tensors and make 2 calls for Assign.
+  // The second Assign call should trigger the case of assign with different
+  // size.
+  TfLiteTensor tensor_a, tensor_b;
+  std::vector<int> shape_a = {1};
+  std::vector<int> shape_b = {2};
+  InitTensor(shape_a, kTfLiteDynamic, 1.0, &tensor_a);
+  InitTensor(shape_b, kTfLiteDynamic, 4.0, &tensor_b);
+
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_a));
+  EXPECT_TRUE(var.IsInitialized());
+  auto* value = var.GetTensor();
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float), value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(1, value->dims->data[0]);
+  EXPECT_EQ(1.0f, value->data.f[0]);
+
+  // Second AssignFrom but now tensor_b has different size from the variable.
+  EXPECT_EQ(kTfLiteOk, var.AssignFrom(&tensor_b));
+  EXPECT_TRUE(var.IsInitialized());
+  value = var.GetTensor();
+  // Variables are always dynamic type.
+  EXPECT_EQ(kTfLiteDynamic, value->allocation_type);
+  EXPECT_EQ(kTfLiteFloat32, value->type);
+  EXPECT_EQ(sizeof(float) * 2, value->bytes);
+  EXPECT_EQ(1, value->dims->size);
+  EXPECT_EQ(2, value->dims->data[0]);
+  EXPECT_EQ(4.0f, value->data.f[0]);
+
+  // Cleanup
+  TfLiteTensorFree(&tensor_a);
+  TfLiteTensorFree(&tensor_b);
+}
+
+}  // namespace resource
+}  // namespace tflite