chore: import upstream snapshot with attribution

tensorflow/lite/experimental/resource/BUILD

@@ -0,0 +1,78 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "cache_buffer",
+    srcs = ["cache_buffer.cc"],
+    hdrs = [
+        "cache_buffer.h",
+        "//tensorflow/lite/core/c:common.h",
+    ],
+    deps = [
+        ":resource",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "cache_buffer_test",
+    srcs = ["cache_buffer_test.cc"],
+    deps = [
+        ":cache_buffer",
+        "//tensorflow/lite/c:common",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "resource",
+    srcs = [
+        "initialization_status.cc",
+        "resource_variable.cc",
+        "static_hashtable.cc",
+    ],
+    hdrs = [
+        "initialization_status.h",
+        "lookup_interfaces.h",
+        "lookup_util.h",
+        "resource_base.h",
+        "resource_variable.h",
+        "static_hashtable.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "//tensorflow/lite/kernels/internal:types",
+    ],
+)
+
+cc_test(
+    name = "resource_variable_test",
+    srcs = [
+        "resource_variable_test.cc",
+    ],
+    deps = [
+        ":resource",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_googletest//:gtest",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/resource/cache_buffer.cc

@@ -0,0 +1,56 @@
+/* Copyright 2024 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/cache_buffer.h"
+
+#include <cstdlib>
+#include <cstring>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace resource {
+
+TfLiteStatus CacheBuffer::Initialize(const TfLiteIntArray& shape) {
+  // Set the dims and allocate the memory.
+  dims_ = TfLiteIntArrayCopy(&shape);
+  const size_t buf_size = NumElements(&shape);
+  buffer_.reset(new float[buf_size]);
+  memset(buffer_.get(), 0, sizeof(float) * buf_size);
+
+  num_entries_.reset(new size_t[shape.data[1]]);
+  memset(num_entries_.get(), 0, sizeof(size_t) * shape.data[1]);
+  is_initialized_ = true;
+  return kTfLiteOk;
+}
+
+size_t CacheBuffer::GetSize() { return sizeof(float) * NumElements(dims_); }
+
+size_t CacheBuffer::GetNumEntries(int idx) const { return num_entries_[idx]; }
+
+CacheBuffer::~CacheBuffer() { TfLiteIntArrayFree(dims_); }
+
+float* CacheBuffer::GetBuffer() { return buffer_.get(); }
+
+void CacheBuffer::SetNumEntries(int idx, size_t count) {
+  TFLITE_DCHECK(count <= dims_->data[2]);
+  num_entries_[idx] = count;
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/cache_buffer.h

@@ -0,0 +1,59 @@
+/* Copyright 2024 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_CACHE_BUFFER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_CACHE_BUFFER_H_
+
+#include <cstddef>
+#include <memory>
+#include <unordered_map>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/resource_variable.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace resource {
+
+/// WARNING: Experimental interface, subject to change.
+// A Cache Buffer class. Useful for keeping the keys and values of a
+// transformer block attention mechanism in autoregressive decode.
+// Ops can access this buffer and add tensors to it. It also keeps track of the
+// number of used entries in the cache.
+class CacheBuffer : public ResourceVariable {
+ public:
+  CacheBuffer() = default;
+  CacheBuffer(const CacheBuffer &) = delete;
+  ~CacheBuffer() override;
+  CacheBuffer &operator=(const CacheBuffer &) = delete;
+  // Initialize tensor of a certain shape using the provided type.
+  TfLiteStatus Initialize(const TfLiteIntArray &shape);
+  size_t GetNumEntries(int idx) const;
+  float *GetBuffer();
+  size_t GetSize();
+  void SetNumEntries(int idx, size_t count);
+
+ private:
+  // The number of entries currently used in the buffer;
+  std::unique_ptr<size_t[]> num_entries_;
+  // The float buffer for storage. Has shape:
+  // <batch, num layers, seq length, num heads, head dim>
+  std::unique_ptr<float[]> buffer_;
+  TfLiteIntArray *dims_;
+};
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_CACHE_BUFFER_H_

tensorflow/lite/experimental/resource/cache_buffer_test.cc

@@ -0,0 +1,44 @@
+/* Copyright 2024 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/cache_buffer.h"
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+namespace resource {
+
+TEST(CacheBufferTest, Initialize) {
+  TfLiteIntArray* shape = TfLiteIntArrayCreate(4);
+  shape->data[0] = 1;
+  shape->data[1] = 3;
+  shape->data[2] = 5;
+  shape->data[3] = 7;
+
+  CacheBuffer cache_buffer;
+  cache_buffer.Initialize(*shape);
+
+  EXPECT_EQ(cache_buffer.GetSize(), 420);
+  ASSERT_NE(cache_buffer.GetBuffer(), nullptr);
+  EXPECT_EQ(cache_buffer.GetNumEntries(0), 0);
+  EXPECT_EQ(cache_buffer.GetNumEntries(1), 0);
+  EXPECT_EQ(cache_buffer.GetNumEntries(2), 0);
+  cache_buffer.SetNumEntries(0, 3);
+  EXPECT_EQ(cache_buffer.GetNumEntries(0), 3);
+  TfLiteIntArrayFree(shape);
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/initialization_status.cc

@@ -0,0 +1,42 @@
+/* Copyright 2021 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/initialization_status.h"
+
+#include <memory>
+
+namespace tflite {
+namespace resource {
+
+void InitializationStatus::MarkInitializationIsDone() {
+  is_initialized_ = true;
+}
+
+bool InitializationStatus::IsInitialized() { return is_initialized_; }
+
+InitializationStatus* GetInitializationStatus(InitializationStatusMap* map,
+                                              int subgraph_id) {
+  auto it = map->find(subgraph_id);
+  if (it != map->end()) {
+    return static_cast<InitializationStatus*>(it->second.get());
+  }
+
+  InitializationStatus* status = new InitializationStatus();
+  map->emplace(subgraph_id, std::unique_ptr<InitializationStatus>(status));
+  return status;
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/initialization_status.h

@@ -0,0 +1,74 @@
+/* Copyright 2021 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_INITIALIZATION_STATUS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_INITIALIZATION_STATUS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <unordered_map>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+
+namespace tflite {
+namespace resource {
+
+/// WARNING: Experimental interface, subject to change.
+// An initialization status class. This class will record the completion status
+// of the initialization procedure. For example, when the initialization
+// subgraph should be invoked  once in a life cycle, this class instance will
+// have the initialization status in order to make sure the followup invocations
+// to invoke the initalization subgraph can be ignored safely.
+class InitializationStatus : public ResourceBase {
+ public:
+  InitializationStatus() {}
+  InitializationStatus(InitializationStatus&& other) noexcept {
+    is_initialized_ = other.is_initialized_;
+  }
+
+  InitializationStatus(const InitializationStatus&) = delete;
+  InitializationStatus& operator=(const InitializationStatus&) = delete;
+
+  ~InitializationStatus() override {}
+
+  ResourceType GetResourceType() const override {
+    return ResourceType::kInitializationStatus;
+  }
+
+  // Mark initialization is done.
+  void MarkInitializationIsDone();
+
+  // Returns true if this initialization is done.
+  bool IsInitialized() override;
+
+  size_t GetMemoryUsage() override { return 0; }
+
+ private:
+  // True if the initialization process is done.
+  bool is_initialized_ = false;
+};
+
+/// WARNING: Experimental interface, subject to change.
+using InitializationStatusMap =
+    std::unordered_map<std::int32_t, std::unique_ptr<InitializationStatus>>;
+
+InitializationStatus* GetInitializationStatus(InitializationStatusMap* map,
+                                              int subgraph_id);
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_INITIALIZATION_STATUS_H_

tensorflow/lite/experimental/resource/lookup_interfaces.h

@@ -0,0 +1,68 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_INTERFACES_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_INTERFACES_H_
+
+#include <unordered_map>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/lookup_util.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace resource {
+
+/// WARNING: Experimental interface, subject to change.
+// A resource hash table interface. It's similar to TensorFlow core's
+// LookupInterface class. But it's identified with int32 ID in TFLite (instead
+// of using Resource handle like TensorFlow).
+class LookupInterface : public ResourceBase {
+ public:
+  ResourceType GetResourceType() const override {
+    return ResourceType::kHashTable;
+  }
+
+  virtual TfLiteStatus Lookup(TfLiteContext* context, const TfLiteTensor* keys,
+                              TfLiteTensor* values,
+                              const TfLiteTensor* default_value) = 0;
+  virtual TfLiteStatus Import(TfLiteContext* context, const TfLiteTensor* keys,
+                              const TfLiteTensor* values) = 0;
+  virtual size_t Size() = 0;
+
+  virtual TfLiteType GetKeyType() const = 0;
+  virtual TfLiteType GetValueType() const = 0;
+  virtual TfLiteStatus CheckKeyAndValueTypes(TfLiteContext* context,
+                                             const TfLiteTensor* keys,
+                                             const TfLiteTensor* values) = 0;
+};
+
+// Creates an resource hash table, shared among all the subgraphs with the
+// given resource id if there is an existing one.
+// WARNING: Experimental interface, subject to change.
+void CreateHashtableResourceIfNotAvailable(ResourceMap* resources,
+                                           int resource_id,
+                                           TfLiteType key_dtype,
+                                           TfLiteType value_dtype);
+
+// Returns the corresponding resource hash table, or nullptr if none.
+// WARNING: Experimental interface, subject to change.
+LookupInterface* GetHashtableResource(ResourceMap* resources, int resource_id);
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_INTERFACES_H_

tensorflow/lite/experimental/resource/lookup_util.h

@@ -0,0 +1,116 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_UTIL_H_
+
+#include <string>
+
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace resource {
+namespace internal {
+
+/// Helper class for accessing TFLite tensor data.
+template <typename T>
+class TensorReader {
+ public:
+  explicit TensorReader(const TfLiteTensor* input) {
+    input_data_ = GetTensorData<T>(input);
+  }
+
+  // Returns the corresponding scalar data at the given index position.
+  // In here, it does not check the validity of the index should be guaranteed
+  // in order not to harm the performance. Caller should take care of it.
+  T GetData(int index) { return input_data_[index]; }
+
+ private:
+  const T* input_data_;
+};
+
+/// Helper class for accessing TFLite tensor data. This specialized class is for
+/// std::string type.
+template <>
+class TensorReader<std::string> {
+ public:
+  explicit TensorReader(const TfLiteTensor* input) : input_(input) {}
+
+  // Returns the corresponding string data at the given index position.
+  // In here, it does not check the validity of the index should be guaranteed
+  // in order not to harm the performance. Caller should take care of it.
+  std::string GetData(int index) {
+    auto string_ref = GetString(input_, index);
+    return std::string(string_ref.str, string_ref.len);
+  }
+
+ private:
+  const TfLiteTensor* input_;
+};
+
+/// WARNING: Experimental interface, subject to change.
+/// Helper class for writing TFLite tensor data.
+template <typename ValueType>
+class TensorWriter {
+ public:
+  explicit TensorWriter(TfLiteTensor* values) {
+    output_data_ = GetTensorData<ValueType>(values);
+  }
+
+  // Sets the given value to the given index position of the tensor storage.
+  // In here, it does not check the validity of the index should be guaranteed
+  // in order not to harm the performance. Caller should take care of it.
+  void SetData(int index, ValueType& value) { output_data_[index] = value; }
+
+  // Commit updates. In this case, it does nothing since the SetData method
+  // writes data directly.
+  void Commit() {
+    // Noop.
+  }
+
+ private:
+  ValueType* output_data_;
+};
+
+/// WARNING: Experimental interface, subject to change.
+/// Helper class for writing TFLite tensor data. This specialized class is for
+/// std::string type.
+template <>
+class TensorWriter<std::string> {
+ public:
+  explicit TensorWriter(TfLiteTensor* values) : values_(values) {}
+
+  // Queues the given string value to the buffer regardless of the provided
+  // index.
+  // In here, it does not check the validity of the index should be guaranteed
+  // in order not to harm the performance. Caller should take care of it.
+  void SetData(int index, const std::string& value) {
+    buf_.AddString(value.data(), value.length());
+  }
+
+  // Commit updates. The stored data in DynamicBuffer will be written into the
+  // tensor storage.
+  void Commit() { buf_.WriteToTensor(values_, nullptr); }
+
+ private:
+  TfLiteTensor* values_;
+  DynamicBuffer buf_;
+};
+
+}  // namespace internal
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_LOOKUP_UTIL_H_

tensorflow/lite/experimental/resource/resource_base.h

@@ -0,0 +1,62 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_BASE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_BASE_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+namespace tflite {
+namespace resource {
+
+// ResourceBase is an abstract base class for resources.
+/// WARNING: Experimental interface, subject to change.
+class ResourceBase {
+ public:
+  enum class ResourceType {
+    kUnknown = 0,
+    kResourceVariable = 1,
+    kHashTable = 2,
+    kInitializationStatus = 3,
+  };
+
+  explicit ResourceBase() {}
+  virtual ~ResourceBase() {}
+
+  virtual ResourceType GetResourceType() const = 0;
+
+  // Returns true if it is initialized.
+  virtual bool IsInitialized() = 0;
+
+  virtual size_t GetMemoryUsage() {
+    return 0;
+  }  // TODO(b/242603814): Make it pure virtual.
+};
+
+/// WARNING: Experimental interface, subject to change.
+using ResourceMap =
+    std::unordered_map<std::int32_t, std::unique_ptr<ResourceBase>>;
+
+using ResourceIDMap = std::map<std::pair<std::string, std::string>, int>;
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_BASE_H_

tensorflow/lite/experimental/resource/resource_variable.cc

@@ -0,0 +1,112 @@
+/* Copyright 2019 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 <cstdlib>
+#include <cstring>
+#include <memory>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+
+namespace tflite {
+namespace resource {
+
+ResourceVariable::ResourceVariable() {
+  memset(&tensor_, 0, sizeof(TfLiteTensor));
+}
+
+ResourceVariable::ResourceVariable(ResourceVariable&& other) {
+  tensor_ = other.tensor_;
+  is_initialized_ = other.is_initialized_;
+
+  memset(&other.tensor_, 0, sizeof(TfLiteTensor));
+  other.is_initialized_ = false;
+}
+
+ResourceVariable::~ResourceVariable() {
+  if (is_initialized_) {
+    free(tensor_.data.raw);
+    if (tensor_.dims) {
+      TfLiteIntArrayFree(tensor_.dims);
+    }
+  }
+}
+
+TfLiteStatus ResourceVariable::AssignFrom(const TfLiteTensor* tensor) {
+  // Save the old allocated resources and attributes that we might use.
+  char* old_raw = tensor_.data.raw;
+  size_t old_bytes = tensor_.bytes;
+  TfLiteIntArray* old_dims = tensor_.dims;
+
+  // Copy primitive parameters.
+  memset(&tensor_, 0, sizeof(tensor_));
+  tensor_.name = "ResourceVariable";
+  tensor_.allocation_type = kTfLiteDynamic;
+  tensor_.type = tensor->type;
+  tensor_.params = tensor->params;
+  tensor_.quantization = tensor->quantization;
+
+  // Copy old shape if possible otherwise create a new one.
+  if (TfLiteIntArrayEqual(old_dims, tensor->dims)) {
+    tensor_.dims = old_dims;
+  } else {
+    TfLiteIntArrayFree(old_dims);
+    tensor_.dims = TfLiteIntArrayCopy(tensor->dims);
+  }
+
+  // Reuse the same buffer if possible otherwise allocate a new one.
+  tensor_.data.raw = old_raw;
+  if (old_bytes != tensor->bytes) {
+    TfLiteTensorRealloc(tensor->bytes, &tensor_);
+  } else {
+    tensor_.bytes = old_bytes;
+  }
+
+  if (tensor->data.raw) {
+    memcpy(tensor_.data.raw, tensor->data.raw, tensor_.bytes);
+  }
+  is_initialized_ = true;
+
+  return kTfLiteOk;
+}
+
+void CreateResourceVariableIfNotAvailable(ResourceMap* resources,
+                                          int resource_id) {
+  if (resources->count(resource_id) != 0) {
+    return;
+  }
+  resources->emplace(resource_id, std::make_unique<ResourceVariable>());
+}
+
+ResourceVariable* GetResourceVariable(ResourceMap* resources, int resource_id) {
+  auto it = resources->find(resource_id);
+  if (it != resources->end() &&
+      it->second->GetResourceType() ==
+          ResourceBase::ResourceType::kResourceVariable) {
+    return static_cast<ResourceVariable*>(it->second.get());
+  }
+  return nullptr;
+}
+
+bool IsBuiltinResource(const TfLiteTensor* tensor) {
+  return tensor && tensor->type == kTfLiteResource &&
+         tensor->delegate == nullptr;
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/resource_variable.h

@@ -0,0 +1,85 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_VARIABLE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_VARIABLE_H_
+
+#include <cstddef>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+
+namespace tflite {
+namespace resource {
+
+/// WARNING: Experimental interface, subject to change.
+// A resource variable class. It's similar to TensorFlow Resource
+// Variable, but it's identified with int32 ID in TFLite (instead of
+// using Resource handle like TensorFlow).
+class ResourceVariable : public ResourceBase {
+ public:
+  ResourceVariable();
+  ResourceVariable(ResourceVariable&& other);
+
+  ResourceVariable(const ResourceVariable&) = delete;
+  ResourceVariable& operator=(const ResourceVariable&) = delete;
+
+  ~ResourceVariable() override;
+
+  ResourceType GetResourceType() const override {
+    return ResourceType::kResourceVariable;
+  }
+
+  // Assigns data from a tensor. Copies its type, shape and data over.
+  TfLiteStatus AssignFrom(const TfLiteTensor* tensor);
+
+  // Get the data tensor stored in the resource variable.
+  // Returns `nullptr` if the variable is never initialized by calling
+  // `AssignFrom`.
+  TfLiteTensor* GetTensor() { return is_initialized_ ? &tensor_ : nullptr; }
+
+  // Returns true if this resource variable is initialized.
+  bool IsInitialized() override { return is_initialized_; }
+
+  size_t GetMemoryUsage() override {
+    return is_initialized_ ? tensor_.bytes : 0;
+  }
+
+ protected:
+  // The tensor (and its buffer stored in `tensor_.data` is fully owned by
+  // the `ResourceVariable` object.
+  TfLiteTensor tensor_;
+  // True if `AssignFrom` function is every called.
+  // False if and only if `tensor_` is filled with zeros.
+  bool is_initialized_ = false;
+};
+
+// Creates a resource variable, shared among all the subgraphs with the given
+// resource id if there is an existing one.
+// WARNING: Experimental interface, subject to change.
+void CreateResourceVariableIfNotAvailable(ResourceMap* resources,
+                                          int resource_id);
+
+// Returns the corresponding resource variable, or nullptr if none.
+// WARNING: Experimental interface, subject to change.
+ResourceVariable* GetResourceVariable(ResourceMap* resources, int resource_id);
+
+// Returns true if 'tensor' points to a builtin resource.
+// WARNING: Experimental interface, subject to change.
+bool IsBuiltinResource(const TfLiteTensor* tensor);
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_RESOURCE_VARIABLE_H_

tensorflow/lite/experimental/resource/resource_variable_test.cc

@@ -0,0 +1,222 @@
+/* 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 <cstdlib>
+#include <cstring>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/test_util.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);
+  ASSERT_THAT(value, DimsAre({1}));
+  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);
+  ASSERT_THAT(value, DimsAre({1}));
+  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);
+  ASSERT_THAT(value, DimsAre({1}));
+  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);
+  ASSERT_THAT(value, DimsAre({1}));
+  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);
+  ASSERT_THAT(value, DimsAre({2}));
+  EXPECT_EQ(4.0f, value->data.f[0]);
+
+  // Cleanup
+  TfLiteTensorFree(&tensor_a);
+  TfLiteTensorFree(&tensor_b);
+}
+
+TEST(IsBuiltinResource, IsBuiltinResourceTest) {
+  TfLiteTensor tensor;
+  tensor.type = kTfLiteResource;
+  tensor.delegate = nullptr;
+  // Resource type and not delegate output.
+  EXPECT_TRUE(IsBuiltinResource(&tensor));
+
+  // Not valid tensor.
+  EXPECT_FALSE(IsBuiltinResource(nullptr));
+
+  // Not a resource type.
+  tensor.type = kTfLiteFloat32;
+  EXPECT_FALSE(IsBuiltinResource(&tensor));
+
+  // Resource but coming from a delegate.
+  tensor.type = kTfLiteResource;
+  TfLiteDelegate delegate;
+  tensor.delegate = &delegate;
+  EXPECT_FALSE(IsBuiltinResource(&tensor));
+}
+
+TEST(ResourceTest, GetMemoryUsage) {
+  ResourceVariable var;
+  EXPECT_FALSE(var.IsInitialized());
+
+  TfLiteTensor tensor;
+  std::vector<int> shape = {100};
+  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(100 * sizeof(float), value->bytes);
+  ASSERT_THAT(value, DimsAre({100}));
+  EXPECT_EQ(1.0f, value->data.f[0]);
+
+  // Check memory usage
+  EXPECT_EQ(100 * sizeof(float), var.GetMemoryUsage());
+
+  // Cleanup
+  // For non dynamic tensors we need to delete the buffers manually.
+  free(tensor.data.raw);
+  TfLiteTensorFree(&tensor);
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/static_hashtable.cc

@@ -0,0 +1,125 @@
+/* Copyright 2019 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/static_hashtable.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/experimental/resource/lookup_interfaces.h"
+#include "tensorflow/lite/experimental/resource/lookup_util.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace resource {
+namespace internal {
+
+template <typename KeyType, typename ValueType>
+TfLiteStatus StaticHashtable<KeyType, ValueType>::Lookup(
+    TfLiteContext* context, const TfLiteTensor* keys, TfLiteTensor* values,
+    const TfLiteTensor* default_value) {
+  if (!is_initialized_) {
+    TF_LITE_KERNEL_LOG(context,
+                       "hashtable need to be initialized before using");
+    return kTfLiteError;
+  }
+  const int size =
+      MatchingFlatSize(GetTensorShape(keys), GetTensorShape(values));
+
+  auto key_tensor_reader = TensorReader<KeyType>(keys);
+  auto value_tensor_writer = TensorWriter<ValueType>(values);
+  auto default_value_tensor_reader = TensorReader<ValueType>(default_value);
+  ValueType first_default_value = default_value_tensor_reader.GetData(0);
+
+  for (int i = 0; i < size; ++i) {
+    auto result = map_.find(key_tensor_reader.GetData(i));
+    if (result != map_.end()) {
+      value_tensor_writer.SetData(i, result->second);
+    } else {
+      value_tensor_writer.SetData(i, first_default_value);
+    }
+  }
+
+  // This is for a string tensor case in order to write buffer back to the
+  // actual tensor destination. Otherwise, it does nothing since the scalar data
+  // will be written into the tensor storage directly.
+  value_tensor_writer.Commit();
+
+  return kTfLiteOk;
+}
+
+template <typename KeyType, typename ValueType>
+TfLiteStatus StaticHashtable<KeyType, ValueType>::Import(
+    TfLiteContext* context, const TfLiteTensor* keys,
+    const TfLiteTensor* values) {
+  // Import nodes can be invoked twice because the converter will not extract
+  // the initializer graph separately from the original graph. The invocations
+  // after the first call will be ignored.
+  if (is_initialized_) {
+    return kTfLiteOk;
+  }
+
+  const int size =
+      MatchingFlatSize(GetTensorShape(keys), GetTensorShape(values));
+
+  auto key_tensor_reader = TensorReader<KeyType>(keys);
+  auto value_tensor_writer = TensorReader<ValueType>(values);
+  for (int i = 0; i < size; ++i) {
+    map_.insert({key_tensor_reader.GetData(i), value_tensor_writer.GetData(i)});
+  }
+
+  is_initialized_ = true;
+  return kTfLiteOk;
+}
+
+LookupInterface* CreateStaticHashtable(TfLiteType key_type,
+                                       TfLiteType value_type) {
+  if (key_type == kTfLiteInt64 && value_type == kTfLiteString) {
+    return new StaticHashtable<std::int64_t, std::string>(key_type, value_type);
+  } else if (key_type == kTfLiteString && value_type == kTfLiteInt64) {
+    return new StaticHashtable<std::string, std::int64_t>(key_type, value_type);
+  }
+  return nullptr;
+}
+
+}  // namespace internal
+
+void CreateHashtableResourceIfNotAvailable(ResourceMap* resources,
+                                           int resource_id,
+                                           TfLiteType key_dtype,
+                                           TfLiteType value_dtype) {
+  if (resources->count(resource_id) != 0) {
+    return;
+  }
+  auto* hashtable = internal::CreateStaticHashtable(key_dtype, value_dtype);
+  resources->emplace(resource_id, std::unique_ptr<LookupInterface>(hashtable));
+}
+
+LookupInterface* GetHashtableResource(ResourceMap* resources, int resource_id) {
+  auto it = resources->find(resource_id);
+  if (it != resources->end() &&
+      it->second->GetResourceType() == ResourceBase::ResourceType::kHashTable) {
+    return static_cast<LookupInterface*>(it->second.get());
+  }
+  return nullptr;
+}
+
+}  // namespace resource
+}  // namespace tflite

tensorflow/lite/experimental/resource/static_hashtable.h

@@ -0,0 +1,87 @@
+/* Copyright 2019 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_STATIC_HASHTABLE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_STATIC_HASHTABLE_H_
+
+#include <cstddef>
+#include <unordered_map>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/resource/lookup_interfaces.h"
+#include "tensorflow/lite/experimental/resource/lookup_util.h"
+#include "tensorflow/lite/experimental/resource/resource_base.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace resource {
+namespace internal {
+
+// A static hash table class. This hash table allows initialization one time in
+// its life cycle. This hash table implements Tensorflow core's HashTableV2 op.
+template <typename KeyType, typename ValueType>
+class StaticHashtable : public tflite::resource::LookupInterface {
+ public:
+  explicit StaticHashtable(TfLiteType key_type, TfLiteType value_type)
+      : key_type_(key_type), value_type_(value_type) {}
+  ~StaticHashtable() override {}
+
+  // Finds the corresponding value of the given keys tensor in the map and
+  // copies the result data to the given values tensor. If there is no matching
+  // value, it will write the default value into the matched position instead.
+  TfLiteStatus Lookup(TfLiteContext* context, const TfLiteTensor* keys,
+                      TfLiteTensor* values,
+                      const TfLiteTensor* default_value) override;
+
+  // Inserts the given key and value tensor data into the hash table.
+  TfLiteStatus Import(TfLiteContext* context, const TfLiteTensor* keys,
+                      const TfLiteTensor* values) override;
+
+  // Returns the item size of the hash table.
+  size_t Size() override { return map_.size(); }
+
+  TfLiteType GetKeyType() const override { return key_type_; }
+  TfLiteType GetValueType() const override { return value_type_; }
+
+  TfLiteStatus CheckKeyAndValueTypes(TfLiteContext* context,
+                                     const TfLiteTensor* keys,
+                                     const TfLiteTensor* values) override {
+    TF_LITE_ENSURE_EQ(context, keys->type, key_type_);
+    TF_LITE_ENSURE_EQ(context, values->type, value_type_);
+    return kTfLiteOk;
+  }
+
+  // Returns true if the hash table is initialized.
+  bool IsInitialized() override { return is_initialized_; }
+
+  size_t GetMemoryUsage() override { return map_.size() * sizeof(ValueType); }
+
+ private:
+  TfLiteType key_type_;
+  TfLiteType value_type_;
+
+  std::unordered_map<KeyType, ValueType> map_;
+  bool is_initialized_ = false;
+};
+
+::tflite::resource::LookupInterface* CreateStaticHashtable(
+    TfLiteType key_type, TfLiteType value_type);
+
+}  // namespace internal
+
+}  // namespace resource
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_RESOURCE_STATIC_HASHTABLE_H_