chore: import upstream snapshot with attribution

tensorflow/lite/kernels/shim/BUILD

@@ -0,0 +1,262 @@
+# This package is a shim layer over TF and TFLite op kernels.
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.bzl", "if_mobile", "if_not_mobile", "tf_cc_test", "tf_copts")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_compatible_with = get_compatible_with_portable(),
+    # The public targets are marked public individually
+    default_visibility = ["//tensorflow:__subpackages__"],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "tf_headers",
+    srcs = [
+        "op_kernel.h",
+        "shape.h",
+        "status_macros.h",
+        "tensor_view.h",
+        "tf_op_shim.h",
+        "tf_tensor_view.h",
+    ],
+    visibility = ["//tensorflow/core:__pkg__"],
+)
+
+cc_library(
+    name = "tf_kernel_shim",
+    deps = [
+        ":op_kernel",
+        ":shape",
+        ":status_macros",
+        ":tensor_view",
+        ":tf_op_shim",
+        ":tf_tensor_view",
+    ],
+)
+
+cc_library(
+    name = "tensor_view",
+    hdrs = ["tensor_view.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/core/platform:tstring",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "shape",
+    srcs = ["shape.cc"],
+    hdrs = ["shape.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "shape_test",
+    srcs = ["shape_test.cc"],
+    deps = [
+        ":shape",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "tf_tensor_view",
+    srcs = ["tf_tensor_view.cc"],
+    hdrs = ["tf_tensor_view.h"],
+    copts = tf_copts(),
+    deps = [
+        ":tensor_view",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+    ]),
+)
+
+tf_cc_test(
+    name = "tf_tensor_view_test",
+    srcs = ["tf_tensor_view_test.cc"],
+    deps = [
+        ":tf_tensor_view",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:tstring",
+    ]),
+)
+
+cc_library(
+    name = "tflite_tensor_view",
+    srcs = ["tflite_tensor_view.cc"],
+    hdrs = ["tflite_tensor_view.h"],
+    deps = [
+        ":tensor_view",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:type_to_tflitetype",
+        "//tensorflow/lite/core/c:common",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/types:variant",
+    ],
+)
+
+cc_test(
+    name = "tflite_tensor_view_test",
+    srcs = ["tflite_tensor_view_test.cc"],
+    deps = [
+        ":test_util",
+        ":tflite_tensor_view",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core:framework_stable",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "op_kernel",
+    hdrs = ["op_kernel.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":shape",
+        ":status_macros",
+        ":tensor_view",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:variant",
+    ],
+)
+
+cc_library(
+    name = "tf_op_shim",
+    srcs = ["tf_op_shim.cc"],
+    hdrs = ["tf_op_shim.h"],
+    copts = tf_copts(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":op_kernel",
+        ":shape",
+        ":status_macros",
+        ":tensor_view",
+        ":tf_tensor_view",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:macros",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+    ]),
+)
+
+cc_library(
+    name = "tflite_op_shim",
+    srcs = ["tflite_op_shim.cc"],
+    hdrs = ["tflite_op_shim.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":op_kernel",
+        ":shape",
+        ":status_macros",
+        ":tensor_view",
+        ":tflite_tensor_view",
+        "//tensorflow/lite:mutable_op_resolver",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "tflite_op_wrapper",
+    hdrs = ["tflite_op_wrapper.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":op_kernel",
+        ":status_macros",
+        "//tensorflow/lite:type_to_tflitetype",
+        "//tensorflow/lite/c:common",
+    ],
+)
+
+cc_test(
+    name = "tflite_op_wrapper_test",
+    srcs = ["tflite_op_wrapper_test.cc"],
+    deps = [
+        ":op_kernel",
+        ":tflite_op_shim",
+        ":tflite_op_wrapper",
+        "//tensorflow/core/platform:tstring",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "status_macros",
+    hdrs = ["status_macros.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "test_util",
+    testonly = True,
+    srcs = ["test_util.cc"],
+    hdrs = ["test_util.h"],
+    deps = [
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:type_to_tflitetype",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "test_util_test",
+    srcs = ["test_util_test.cc"],
+    deps = [
+        ":test_util",
+        "//tensorflow/lite:cc_api_experimental",
+        "//tensorflow/lite/core:framework_stable",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/kernels/shim/README.md

@@ -0,0 +1,174 @@
+This folder contains a convenience library called *tf-shim* over TF and TFLite
+op kernel APIs.
+
+## Summary
+
+This library creates a shim over the custom op APIs of TF and TFLite so the
+developer can write the custom op once with minimal binary or runtime overhead.
+
+An example usage is an input preprocessing op kernel that can be used in
+both TF and TFLite.
+
+## Background
+
+When there is a need to implement a logic that is not supported by the TF
+builtin ops the alternative is to build a custom op. If that op needs to
+run on-device then it needs to be written in C++ against the client API for
+custom ops.
+
+For example, feature processing especially for textual input in an ML model
+can involve operations that don't lend themselves well to vectorization and the
+code, if written as a C++ function, would be much shorter and more readable.
+
+However, Tensorflow and TFLite APIs for creating op kernels are, at the moment,
+not identical. This library offers a convenient way to write the kernel once and
+adapt it to both TF and TFLite with minimal binary and runtime overhead.
+
+## Implementation
+
+This folder contains two pieces:
+
+1.  `TensorView` as a shim over `::tensorflow::Tensor` and `TfLiteTensor`
+
+2.  `OpKernelShim` class which abstracts the TF and TFLite op kernel APIs.
+
+### TensorView
+
+This class is a *view* over an already allocated tensor in TF or TFLite without
+taking any ownership. In that sense it is similar to `absl::string_view` but with
+the difference that the underlying buffer can be mutable.
+
+Example Usage:
+
+```
+::tensorflow::Tensor tf_tensor;
+auto t = TensorView::New(&tf_tensor);
+
+auto t_str_mat = t.As<::tensorflow::tstring, /*RANK=*/ 2>();
+t(0, 0) = "ab";
+t(0, 1) = "cde"
+
+
+auto t_buffer = t.Data<::tensorflow::tstring>();
+t[0] = "ab";
+t[1] = "cde"
+```
+
+```
+TfLiteTensor tflite_tensor;
+auto t = TensorView::New(&tflite_tensor);
+
+auto t_int_vec = t.As<int32, /*RANK=*/ 1>();
+t(0) = 123;
+t(1) = 456
+
+auto t_buffer = t.Data<int32>();
+t[0] = 123;
+t[1] = 456
+```
+
+The `New` is the factory function which based on the type of the input returns
+either a `TfTensorView` or a `TfLiteTensorView`.
+
+See the unit tests `tf_tensor_view_test.cc` and `tflite_tensor_view_test.cc` for
+more usage.
+
+The string tensor in `TfLiteTensorView` is a bit of special case. Since string
+tensors in TfLite are serialized in a specific format, while writing to those
+tensors an intermediate buffer is needed to hold intermediate values before all
+the strings get serialized. The intermediate string buffers are serialized back
+to the TfLite string format once the last remaining `TfLiteTensorView` goes out
+of scope. Only then the user can see the string values in the underlying
+`TfLiteTensor`. That said, when implementing an op kernel, there is rarely a
+need to read back the contents of a mutable output `TfLiteTensor` within the
+same code block.
+
+### OpKernelShim
+
+*WARNING: Experimental interface, subject to change*
+
+This class defines the interface which when implemented allows for convenient
+adaptation to TF and TFLite op kernels.
+
+Here is an example op kernel implementing this interface:
+
+```
+template<TfRuntime R>
+class MyOp : public OpKernelShim<MyOp, R> {
+
+  // Attributes declaration (syntax: https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Attrs();
+
+  // Input tensors declaration (syntax: https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Inputs();
+
+  // Output tensors declaration (syntax: https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Outputs();
+
+  // Initializes the op
+  absl::Status Init(InitContext* ctx);
+
+  // Runs the operation
+  absl::Status Invoke(InvokeContext* ctx);
+
+  // Shape inference
+  static absl::Status ShapeInference(ShapeInferenceContext* ctx);
+};
+```
+
+The class `MyOp` is passing itself to `OpKernelShim` as a template parameter.
+This is because `OpKernelShim` is a static interface using the CRTP pattern.
+Similarly, the context classes: `InitContext`, `InvokeContext` and
+`ShapeInferenceContext` are all static interfaces in the same way.
+
+The class `MyOp` can also be templatized. See `test_op/tmpl_op.h` for an
+example.
+
+### Context Interfaces
+
+An op kernel written using this library has access to a number of *context*
+objects at various stages of its lifecycle. These context objects are
+effectively shims over the existing context objects in TF and TFLite.
+
+#### InitContext
+An instance of this class is passed to the op kernel during its initialization.
+
+```
+template <typename SubType>
+class InitContext {
+ public:
+  // Read the given attribute and populate the given value.
+  template <typename AttrType>
+  absl::Status GetAttr(const std::string& attr_name, AttrType* value) const;
+};
+```
+
+#### InvokeContext
+An instance of this class is passed to the op kernel during its invocation.
+
+```
+template <typename SubType>
+class InvokeContext {
+ public:
+  // Read an input tensor
+  ConstTensorViewOr GetInput(const int idx) const;
+  // Get a mutable output tensor
+  TensorViewOr GetOutput(const int idx, const Shape& shape) const;
+};
+```
+
+#### ShapeInferenceContext
+An instance of this class is passed to the op kernel during its shape inference.
+
+```
+template <typename SubType>
+class ShapeInferenceContext {
+ public:
+  // Read an input tensor shape
+  ShapeOr GetInputShape(const int idx) const;
+  // Set an output tensor shape
+  absl::Status SetOutputShape(const int idx, const Shape& shape);
+  // Read an input tensor during shape inference
+  ConstTensorViewOr GetInputTensor(const int idx) const;
+};
+```

tensorflow/lite/kernels/shim/op_kernel.h

@@ -0,0 +1,279 @@
+/* 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_KERNELS_SHIM_OP_KERNEL_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_OP_KERNEL_H_
+
+// This file defines a shim layer on top of TF and TFLite custom op APIs.
+// The goal is for a custom op to be written once and used for both runtimes
+//
+// It consists of two pieces:
+//   * A set of *context* interfaces:
+//     ** InvokeContext, InitContext, ShapeInferenceContext
+//     These are passed on to the custom op implementation to read/write
+//     tensors, etc.
+//
+//   * An OpKernelShim interface:
+//     This is what a custom op needs to implement. By using that interface the
+//     custom op can then be easily converted to both a TF op kernel and a
+//     TFLite op kernel.
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <variant>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/variant.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+
+namespace tflite {
+namespace shim {
+
+// List of the TF custom op APIs this shim library is abstracting away.
+//
+// This enum is used as the template parameter in various places in
+// order to pick the correct set of types (eg. TfInvokeContext vs.
+// TfLiteInvokeContext) in the op implementation.
+enum class Runtime { kTf, kTfLite };
+
+// TensorView or error
+using TensorViewOr = absl::StatusOr<std::unique_ptr<TensorView>>;
+using ConstTensorViewOr = absl::StatusOr<std::unique_ptr<const TensorView>>;
+
+// Below are the interfaces for various "Context" objects to abstract away the
+// TF and TFLite differences.
+//
+// The interfaces are static and use the CRTP pattern instead of virtual
+// methods.
+
+// The attribute dictionary passed to the op
+using AttrValue = std::variant<bool, int64_t, float, absl::string_view>;
+
+// The interface for available methods during an op kernel initialization
+template <typename SubType>
+class InitContext {
+ public:
+  // Read the given attribute and populate the given value.
+  template <typename AttrType>
+  absl::Status GetAttr(const std::string& attr_name, AttrType* value) const;
+
+ protected:
+  // Read a given attribute or return error
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const {
+    return static_cast<const SubType&>(*this).GetAttr(attr_name);
+  }
+};
+
+// The interface for available methods during an op kernel invocation
+template <typename SubType>
+class InvokeContext {
+ public:
+  // Read an input tensor
+  ConstTensorViewOr GetInput(const int idx) const {
+    return static_cast<const SubType&>(*this).GetInput(idx);
+  }
+  // Get a mutable output tensor
+  TensorViewOr GetOutput(const int idx, const Shape& shape) const {
+    return static_cast<const SubType&>(*this).GetOutput(idx, shape);
+  }
+  // Number of input tensors
+  int NumInputs() const {
+    return static_cast<const SubType&>(*this).NumInputs();
+  }
+  // Number of output tensors
+  int NumOutputs() const {
+    return static_cast<const SubType&>(*this).NumOutputs();
+  }
+};
+
+// The interface for available methods during shape inference
+template <typename SubType>
+class ShapeInferenceContext {
+ public:
+  // Read an input tensor shape
+  ShapeOr GetInputShape(const int idx) const {
+    return static_cast<const SubType&>(*this).GetInputShape(idx);
+  }
+  // Set an output tensor shape
+  absl::Status SetOutputShape(const int idx, const Shape& shape) {
+    return static_cast<SubType&>(*this).SetOutputShape(idx, shape);
+  }
+  // Read an input tensor during shape inference
+  ConstTensorViewOr GetInputTensor(const int idx) const {
+    return static_cast<const SubType&>(*this).GetInputTensor(idx);
+  }
+  // Number of input tensors
+  int NumInputs() const {
+    return static_cast<const SubType&>(*this).NumInputs();
+  }
+  // Number of output tensors
+  int NumOutputs() const {
+    return static_cast<const SubType&>(*this).NumOutputs();
+  }
+  // Read the given attribute and populate the given value.
+  template <typename AttrType>
+  absl::Status GetAttr(const std::string& attr_name, AttrType* value) const;
+
+ protected:
+  // Read a given attribute or return error
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const {
+    return static_cast<const SubType&>(*this).GetAttr(attr_name);
+  }
+};
+
+// Maps the Runtime to the correct context types.
+// eg. ContextTypeForRuntime<Runtime::Tf>  -->
+//       { TfInitContext, TfInvokeContext, TfShapreInferenceContext }
+template <Runtime Rt>
+struct ContextTypeForRuntime {
+  // * Init
+  // * Invoke
+  // * ShapeInference
+};
+
+// A Tensorflow operation interface which is then adapted to both TF and TFLite
+// runtimes.
+//
+// Example usage:
+//
+//   template<Runtime R>
+//   class MyOp : public OpKernelShim<MyOp, R> {
+//
+//     // Attributes declaration
+//     // (syntax: https://www.tensorflow.org/guide/create_op)
+//     static std::vector<std::string> Attrs();
+//
+//     // Input tensors declaration
+//     // (syntax: https://www.tensorflow.org/guide/create_op)
+//     static std::vector<std::string> Inputs();
+//
+//     // Output tensors declaration
+//     // (syntax: https://www.tensorflow.org/guide/create_op)
+//     static std::vector<std::string> Outputs();
+//
+//     // Initializes the op
+//     absl::Status Init(InitContext* ctx);
+//
+//     // Runs the operation
+//     absl::Status Invoke(InvokeContext* ctx);
+//
+//     // Shape inference
+//     static absl::Status ShapeInference(ShapeInferenceContext* ctx);
+//
+//   };
+//
+// WARNING: Experimental interface, subject to change
+template <template <Runtime, typename...> typename SubType, Runtime Rt,
+          typename... Ts>
+class OpKernelShim {
+ public:
+  // Some typedefs for convenience
+  using Shape = ::tflite::shim::Shape;
+  using InitContext =
+      ::tflite::shim::InitContext<typename ContextTypeForRuntime<Rt>::Init>;
+  using InvokeContext =
+      ::tflite::shim::InvokeContext<typename ContextTypeForRuntime<Rt>::Invoke>;
+  using ShapeInferenceContext = ::tflite::shim::ShapeInferenceContext<
+      typename ContextTypeForRuntime<Rt>::ShapeInference>;
+
+  // Needed because the pointer to this class is stored
+  virtual ~OpKernelShim() = default;
+
+  // If the operation has any attributes they are passed here.
+  absl::Status Init(InitContext* ctx) {
+    return static_cast<SubType<Rt, Ts...>&>(*this).Init(ctx);
+  }
+
+  // The actual computations of the operation
+  absl::Status Invoke(InvokeContext* ctx) {
+    return static_cast<SubType<Rt, Ts...>&>(*this).Invoke(ctx);
+  }
+
+  // Shape inference
+  static absl::Status ShapeInference(ShapeInferenceContext* ctx) {
+    return SubType<Rt, Ts...>::ShapeInference(ctx);
+  }
+
+ protected:
+  OpKernelShim() = default;
+
+  // Convience method for filling a single dimension output tensor.
+  template <typename BufferType, typename DType>
+  absl::Status FillOutputTensor(const std::vector<BufferType>& buffer,
+                                int index, InvokeContext* context) const;
+};
+
+/////////////////////// Implementations
+
+namespace internal {
+// Extract the given AttrType from the AttrValue variant or returns error.
+template <typename AttrType>
+absl::Status GetAttr(const std::string& attr_name,
+                     const absl::StatusOr<AttrValue> attr_value_or,
+                     AttrType* value) {
+  if (!attr_value_or.ok()) return attr_value_or.status();
+  const AttrValue& attr_value = attr_value_or.value();
+  if (!std::holds_alternative<AttrType>(attr_value)) {
+    return absl::InternalError(
+        absl::StrCat("The attribute type does not match the provided "
+                     "type: attr_name: ",
+                     attr_name));
+  }
+  *value = std::get<AttrType>(attr_value);
+  return absl::OkStatus();
+}
+}  // namespace internal
+
+template <typename SubType>
+template <typename AttrType>
+absl::Status InitContext<SubType>::GetAttr(const std::string& attr_name,
+                                           AttrType* value) const {
+  const auto attr_value_or = GetAttr(attr_name);
+  return internal::GetAttr<AttrType>(attr_name, attr_value_or, value);
+}
+
+template <typename SubType>
+template <typename AttrType>
+absl::Status ShapeInferenceContext<SubType>::GetAttr(
+    const std::string& attr_name, AttrType* value) const {
+  const auto attr_value_or = GetAttr(attr_name);
+  return internal::GetAttr<AttrType>(attr_name, attr_value_or, value);
+}
+
+template <template <Runtime, typename...> typename SubType, Runtime Rt,
+          typename... Ts>
+template <typename BufferType, typename DType>
+absl::Status OpKernelShim<SubType, Rt, Ts...>::FillOutputTensor(
+    const std::vector<BufferType>& buffer, const int index,
+    tflite::shim::InvokeContext<typename ContextTypeForRuntime<Rt>::Invoke>*
+        context) const {
+  SH_ASSIGN_OR_RETURN(
+      const auto tensorview,
+      context->GetOutput(
+          index, tflite::shim::Shape({static_cast<int>(buffer.size())})));
+  auto data = tensorview->template As<DType, 1>();
+  for (int i = 0; i < buffer.size(); ++i) data(i) = buffer.at(i);
+  return absl::OkStatus();
+}
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_OP_KERNEL_H_

tensorflow/lite/kernels/shim/shape.cc

@@ -0,0 +1,84 @@
+/* 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/kernels/shim/shape.h"
+
+#include <string>
+
+namespace tflite {
+namespace shim {
+
+bool Shape::operator==(const Shape& rhs) const {
+  if (!has_value() || !rhs.has_value()) return false;
+  if (value_.size() != rhs.value_.size()) return false;
+  for (int i = 0; i < value_.size(); ++i)
+    if (value_[i] != rhs.value_[i] || value_[i] == kUnknownDim ||
+        rhs.value_[i] == kUnknownDim)
+      return false;
+  return true;
+}
+
+bool Shape::operator!=(const Shape& rhs) const { return !(*this == rhs); }
+
+bool Shape::Compatible(const Shape& rhs) const {
+  if (!has_value() || !rhs.has_value()) return true;
+  if (value_.size() != rhs.value_.size()) return false;
+  for (int i = 0; i < value_.size(); ++i) {
+    const auto lhs_i = value_[i];
+    const auto rhs_i = rhs.value_[i];
+    if (lhs_i != rhs_i && lhs_i != kUnknownDim && rhs_i != kUnknownDim)
+      return false;
+  }
+  return true;
+}
+
+std::string Shape::ToString() const {
+  std::string ret;
+  if (has_value()) {
+    ret += "[";
+    if (!value_.empty()) ret += " ";
+    for (const auto dim : value_) {
+      if (dim != kUnknownDim) {
+        ret += std::to_string(dim);
+      } else {
+        ret += "?";
+      }
+      ret += " ";
+    }
+    ret += "]";
+  } else {
+    ret += "?";
+  }
+  return ret;
+}
+
+bool Shape::FullyDefined() const {
+  if (!has_value_) return false;
+  for (const auto dim : value_)
+    if (dim == kUnknownDim) return false;
+  return true;
+}
+
+int Shape::AddDims(const int dim1, const int dim2) {
+  if (dim1 == kUnknownDim || dim2 == kUnknownDim) return kUnknownDim;
+  return dim1 + dim2;
+}
+
+int Shape::Dim(const int idx) const {
+  if (has_value_) return value_[idx];
+  return kUnknownDim;
+}
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/shape.h

@@ -0,0 +1,96 @@
+/* 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_KERNELS_SHIM_SHAPE_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_SHAPE_H_
+
+#include <initializer_list>
+#include <string>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+
+namespace tflite {
+namespace shim {
+
+// Shape of a tensor. When unset it means the rank is unknown. Individual dims
+// can also be unknown.
+class Shape {
+ public:
+  using ValueType = std::vector<int>;
+
+  Shape() = default;
+  Shape(const Shape& o) = default;
+  Shape(Shape&& o) = default;
+  Shape& operator=(const Shape& o) = default;
+  Shape& operator=(Shape&& o) = default;
+
+  // Ctors
+  Shape(const std::initializer_list<int>& o) : value_(o), has_value_(true) {}
+  template <typename... Args>
+  explicit Shape(Args&&... args)  // forward ctor args to that of std::vector
+      : value_(std::forward<Args>(args)...), has_value_(true) {}
+  explicit Shape(const absl::Span<int> value)
+      : value_(value.data(), value.data() + value.size()), has_value_(true) {}
+
+  // Accessors
+  inline bool has_value() const { return has_value_; }
+  inline ValueType& value() { return value_; }
+  inline const ValueType& value() const { return value_; }
+  ValueType* operator->() { return &value_; }
+  const ValueType* operator->() const { return &value_; }
+  ValueType& operator*() { return value_; }
+  const ValueType& operator*() const { return value_; }
+  // Get the specified dimension if known
+  int Dim(const int idx) const;
+
+  // Returns the rank of the shape
+  int Rank() const { return has_value_ ? value_.size() : kUnknownRank; }
+
+  // Whether all the dimensions of the shape are known
+  bool FullyDefined() const;
+
+  // Pretty printer
+  std::string ToString() const;
+
+  // Adds two dimension taking into account unknown dims.
+  static int AddDims(const int dim1, const int dim2);
+
+  // Comparison
+
+  // Strict equality of the shapes. Unknown dims or rank on one side will
+  // result in false
+  bool operator==(const Shape& rhs) const;
+  bool operator!=(const Shape& rhs) const;
+
+  // Compatibility of the shapes. If there are two known and incompatible
+  // dimensions it returns false
+  bool Compatible(const Shape& rhs) const;
+
+  // The value for unknown dimensions and rank. There are static_asserts to
+  // ensure this matches the one defined in ::tensorflow namespace
+  static constexpr int kUnknownDim = -1;
+  static constexpr int kUnknownRank = -1;
+
+ private:
+  ValueType value_;
+  bool has_value_ = false;
+};
+using ShapeOr = absl::StatusOr<Shape>;
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_SHAPE_H_

tensorflow/lite/kernels/shim/shape_test.cc

@@ -0,0 +1,104 @@
+/* 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/kernels/shim/shape.h"
+
+#include <gtest/gtest.h>
+
+namespace tflite {
+namespace shim {
+namespace {
+
+TEST(Shape, Eq) {
+  EXPECT_TRUE(Shape({1, 2}) == Shape({1, 2}));
+  // scalar
+  EXPECT_TRUE(Shape(std::vector<int>{}) == Shape(std::vector<int>{}));
+  // vector of size 1
+  EXPECT_TRUE(Shape({1}) == Shape({1}));
+  // dim mismatch
+  EXPECT_FALSE(Shape({1, 2, 1}) == Shape({1, 2}));
+  // dim mismatch
+  EXPECT_FALSE(Shape({1, 3}) == Shape({1, 2}));
+  // unknown dim
+  EXPECT_FALSE(Shape({3, -1, 2}) == Shape({3, -1, 2}));
+  // two unknown ranks
+  EXPECT_FALSE(Shape() == Shape());
+  // Unknown rank vs known shape
+  EXPECT_FALSE(Shape({3, 4}) == Shape());
+  // scalar vs. vector of size 1
+  EXPECT_FALSE(Shape(std::vector<int>{}) == Shape({1}));
+  // unknown dim
+  EXPECT_FALSE(Shape({1, -1}) == Shape({1, 2}));
+  // same num elems
+  EXPECT_FALSE(Shape({1, 2}) == Shape({2, 1}));
+}
+
+TEST(Shape, Compatible) {
+  EXPECT_TRUE(Shape({1, 2}).Compatible(Shape({1, 2})));
+  // scalar
+  EXPECT_TRUE(Shape(std::vector<int>{}).Compatible(Shape(std::vector<int>{})));
+  // vector of size 1
+  EXPECT_TRUE(Shape({1}).Compatible(Shape({1})));
+  // unknown dim
+  EXPECT_TRUE(Shape({3, -1, 2}).Compatible(Shape({3, -1, 2})));
+  // dim mismatch
+  EXPECT_FALSE(Shape({1, 2, 1}).Compatible(Shape({1, 2})));
+  // dim mismatch
+  EXPECT_FALSE(Shape({1, 3}).Compatible(Shape({1, 2})));
+  // two unknown ranks
+  EXPECT_TRUE(Shape().Compatible(Shape()));
+  // Unknown rank vs known shape
+  EXPECT_TRUE(Shape({3, 4}).Compatible(Shape()));
+  // scalar vs. vector of size 1
+  EXPECT_FALSE(Shape(std::vector<int>{}).Compatible(Shape({1})));
+  // unknown dim
+  EXPECT_TRUE(Shape({1, -1}).Compatible(Shape({1, 2})));
+  // same num elems
+  EXPECT_FALSE(Shape({1, 2}).Compatible(Shape({2, 1})));
+}
+
+TEST(Shape, ToStr) {
+  EXPECT_EQ("[ 1 2 ]", Shape({1, 2}).ToString());
+  EXPECT_EQ("[]", Shape({}).ToString());
+  EXPECT_EQ("?", Shape().ToString());
+  EXPECT_EQ("[ 1 2 3 ]", Shape({1, 2, 3}).ToString());
+  EXPECT_EQ("[ 2 ? 3 ]", Shape({2, -1, 3}).ToString());
+}
+
+TEST(Shape, FullyDefined) {
+  EXPECT_TRUE(Shape({}).FullyDefined());
+  EXPECT_FALSE(Shape().FullyDefined());
+  EXPECT_TRUE(Shape({1}).FullyDefined());
+  EXPECT_TRUE(Shape({1, 2}).FullyDefined());
+  EXPECT_FALSE(Shape({3, -1, 2}).FullyDefined());
+  EXPECT_FALSE(Shape({-1, -1}).FullyDefined());
+}
+
+TEST(Shape, Dim) {
+  EXPECT_EQ(-1, Shape().Dim(2));
+  EXPECT_EQ(3, Shape({2, 3}).Dim(1));
+  EXPECT_EQ(3, Shape({-1, 3}).Dim(1));
+  EXPECT_EQ(-1, Shape({-1, 3}).Dim(0));
+}
+
+TEST(Shape, AddDims) {
+  EXPECT_EQ(5, Shape::AddDims(3, 2));
+  EXPECT_EQ(-1, Shape::AddDims(-1, 2));
+  EXPECT_EQ(-1, Shape::AddDims(1, -2));
+  EXPECT_EQ(-1, Shape::AddDims(-1, -1));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/status_macros.h

@@ -0,0 +1,73 @@
+/* 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_KERNELS_SHIM_STATUS_MACROS_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_STATUS_MACROS_H_
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+
+// A few macros to help with propagating status.
+// They are mainly adapted from the ones in  tensorflow/core/platform/errors.h
+// and tensorflow/core/platform/statusor.h but these files come with many
+// transitive deps which can be too much for TFLite use cases. If these type of
+// macros end up in `absl` we can replace them with those.
+
+// The macros are prefixed with SH_ to avoid name collision.
+
+namespace tflite {
+namespace shim {
+template <typename... Args>
+void AppendToMessage(::absl::Status* status, Args... args) {
+  *status = ::absl::Status(status->code(),
+                           ::absl::StrCat(status->message(), "\n\t", args...));
+}
+}  // namespace shim
+}  // namespace tflite
+
+// Propagates error up the stack and appends to the error message.
+#define SH_RETURN_WITH_CONTEXT_IF_ERROR(expr, ...)            \
+  do {                                                        \
+    ::absl::Status _status = (expr);                          \
+    if (!_status.ok()) {                                      \
+      ::tflite::shim::AppendToMessage(&_status, __VA_ARGS__); \
+      return _status;                                         \
+    }                                                         \
+  } while (0)
+
+// Propages the error up the stack.
+// This can't be merged with the SH_RETURN_WITH_CONTEXT_IF_ERROR unless some
+// overly clever/unreadable macro magic is used.
+#define SH_RETURN_IF_ERROR(...)             \
+  do {                                      \
+    ::absl::Status _status = (__VA_ARGS__); \
+    if (!_status.ok()) return _status;      \
+  } while (0)
+
+// Internal helper for concatenating macro values.
+#define SH_STATUS_MACROS_CONCAT_NAME_INNER(x, y) x##y
+#define SH_STATUS_MACROS_CONCAT_NAME(x, y) \
+  SH_STATUS_MACROS_CONCAT_NAME_INNER(x, y)
+
+// Assigns an expression to lhs or propagates the error up.
+#define SH_ASSIGN_OR_RETURN(lhs, rexpr) \
+  SH_ASSIGN_OR_RETURN_IMPL(             \
+      SH_STATUS_MACROS_CONCAT_NAME(statusor, __COUNTER__), lhs, rexpr)
+
+#define SH_ASSIGN_OR_RETURN_IMPL(statusor, lhs, rexpr) \
+  auto statusor = (rexpr);                             \
+  if (!statusor.ok()) return statusor.status();        \
+  lhs = std::move(statusor.value())
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_STATUS_MACROS_H_

tensorflow/lite/kernels/shim/tensor_view.h

@@ -0,0 +1,256 @@
+/* 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_KERNELS_SHIM_TENSOR_VIEW_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TENSOR_VIEW_H_
+
+#include <variant>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/tstring.h"
+
+namespace tflite {
+namespace shim {
+
+// A type deduction template which is specialized for TF and TFLite.
+// That is it maps
+//   ::tensorflow::Tensor -> tflite::shim::TfTensorView
+//   ::TfLiteTensor -> tflite::shim::TfLiteTensorView
+template <typename W>
+struct TensorViewSubType {};
+
+// Common denominator for ::tflite::TfLiteTensor and ::tensorflow::Tensor.
+// It is a "view" over the underlying tensor without taking ownership.
+// Objects of this class can also mutate the underlying tensor depending on
+// whether the underlying tensor is "const" qualified or not.
+//
+// Movable and copyable.
+// It can be instantiated with the New() factory function. eg.
+//   TfTensorView t           = TensorView::New(&tf_tensor);
+//   const TfTensorView t     = TensorView::New(&const_tf_tensor);
+//   TfLiteTensorView t       = TensorView::New(&tflite_tensor);
+//   const TfLiteTensorView t = TensorView::New(&const_tflite_tensor);
+class TensorView {
+ protected:
+  // Union over all data types
+  using DataVariantType =
+      std::variant<absl::Span<bool>, absl::Span<int8_t>, absl::Span<uint8_t>,
+                   absl::Span<int16_t>, absl::Span<uint16_t>,
+                   absl::Span<int32_t>, absl::Span<uint32_t>,
+                   absl::Span<int64_t>, absl::Span<uint64_t>, absl::Span<float>,
+                   absl::Span<double>, absl::Span<::tensorflow::tstring>>;
+
+  // An interface while provides convenient row-major indexing over the
+  // underlying tensor.
+  // Example usage:
+  //
+  //   // A scalar view
+  //   const TensorView t_float
+  //   float val = t_float.AsScalar<float>();
+  //
+  //   // A vector view
+  //   const TensorView t_int;
+  //   auto t_int_vec = t_int.As<int32_t, /*RANK=*/ 1>();
+  //   int sum = t_int_vec(0) + t_int_vec(1);
+  //
+  //   // A matrix view
+  //   TensorView t_str;
+  //   auto t_str_mat = t_str.As<tensorflow::tstring, /*RANK=*/ 2>();
+  //   t_str_mat(0, 0) = "abc";
+  //   t_str_mat(2, 3) = "def";
+  template <typename DType, int RANK>
+  class Tensor {
+   public:
+    explicit Tensor(TensorView *t)
+        : data_(t->Data<DType>()), shape_(t->Shape()) {
+      DCHECK_EQ(RANK, shape_.size());
+      ComputeRowSizes();
+    }
+
+    explicit Tensor(const TensorView *t)
+        : data_(t->Data<DType>()), shape_(t->Shape()) {
+      DCHECK_EQ(RANK, shape_.size());
+      ComputeRowSizes();
+    }
+
+    // indexing operator
+    template <typename... IndexTypes>
+    inline DType &operator()(IndexTypes... indices) {
+      const auto idx = RowMajorIndex(std::array<int, RANK>{{indices...}});
+      return data_[idx];
+    }
+
+    // const indexing operator
+    template <typename... IndexTypes>
+    inline const DType &operator()(IndexTypes... indices) const {
+      const auto idx = RowMajorIndex(std::array<int, RANK>{{indices...}});
+      return data_.at(idx);
+    }
+
+    // Pointer accessor
+    typename absl::Span<DType>::pointer Ptr() { return data_.data(); }
+    constexpr typename absl::Span<DType>::const_pointer Ptr() const {
+      return data_.data();
+    }
+
+    // Size of the given dimension
+    inline int Dim(int dim_i) const {
+      DCHECK(RANK > 0 && dim_i < RANK) << "dim: " << dim_i << " rank:" << RANK;
+      // Handle negative indices
+      if (dim_i < 0) dim_i = ((dim_i % RANK) + RANK) % RANK;
+      return shape_[dim_i];
+    }
+
+    // The tensor's rank: number of dimensions
+    /*[[nodiscard]]*/ constexpr std::size_t Rank() const { return RANK; }
+
+   private:
+    // Computes the row-major index
+    inline std::size_t RowMajorIndex(
+        const std::array<int, RANK> &indices) const {
+      std::size_t ret = 0;
+      for (int i = 0; i < RANK; ++i) ret += indices[i] * row_sizes_[i];
+      return ret;
+    }
+
+    // Pre computes row sizes to convert multi dim indices into a row major
+    // index
+    void ComputeRowSizes() {
+      // Precompute row sizes for row major index computation
+      if (RANK > 0) {
+        row_sizes_[RANK - 1] = 1;
+        for (int i = RANK - 2; i >= 0; --i) {
+          row_sizes_[i] = row_sizes_[i + 1] * shape_[i + 1];
+        }
+      }
+    }
+
+    absl::Span<DType> data_;
+    const absl::Span<int> shape_;
+    std::size_t row_sizes_[RANK]{};
+  };
+
+ public:
+  // Factory which gets specialized for different wrapped tensor types.
+  template <typename W>
+  static absl::StatusOr<typename TensorViewSubType<W>::Type> New(
+      W *wrapped_tensor);
+
+ protected:
+  // Move constructor
+  TensorView(TensorView &&o) = default;
+  // Copy constructor
+  TensorView(const TensorView &o) = default;
+  // Move assignment operator
+  TensorView &operator=(TensorView &&o) = default;
+  // Copy assignment operator
+  TensorView &operator=(const TensorView &) = default;
+
+ public:
+  // Dtor
+  virtual ~TensorView() = default;
+
+  // Accessors
+
+  // Shape
+  absl::Span<int> Shape() { return shape_; }
+  /*[[nodiscard]]*/ const absl::Span<int> Shape() const { return shape_; }
+
+  // Data
+  template <typename DType>
+  absl::Span<DType> &Data() {
+    return std::get<absl::Span<DType>>(data_);
+  }
+  template <typename DType>
+  constexpr absl::Span<DType> Data() const {
+    return std::get<absl::Span<DType>>(data_);
+  }
+
+  // Reads the tensor given the dtype and its rank and provides an indexing
+  // operator.
+  template <typename DType, int RANK>
+  Tensor<DType, RANK> As() {
+    return Tensor<DType, RANK>(this);
+  }
+
+  // Const version of As()
+  template <typename DType, int RANK>
+  const Tensor<DType, RANK> As() const {
+    return Tensor<DType, RANK>(this);
+  }
+
+  // Read the given tensor as a scalar or return error if it isn't
+  template <typename DType>
+  DType &AsScalar();
+
+  template <typename DType>
+  const DType &AsScalar() const;
+
+ protected:
+  // Templated constructor. Since it's not possible to specify the template
+  // argument directly we place a dummy argument of that type so compiler
+  // can deduce the right template parameter
+  template <typename DType>
+  TensorView(const absl::Span<int> shape, void *data,
+             const std::size_t data_size, const DType &)
+      : shape_(shape),
+        data_(absl::Span<DType>(reinterpret_cast<DType *>(data),
+                                data_size / sizeof(DType))) {}
+
+  // Return the total number of elements given the shape.
+  static constexpr std::size_t NumElements(const absl::Span<int> shape) {
+    std::size_t ret = 1;
+    for (const auto dim : shape) ret *= dim;
+    return ret;
+  }
+
+  // Tensor shape
+  // Note: using int rather than size_t to avoid conversion to from TfLite shape
+  absl::Span<int> shape_;
+  // Tensor data
+  DataVariantType data_;
+};
+
+// Add or remove const qualifier to O based on whether it is in I.
+// For example
+//   MatchConstNess<const TfLiteTensor, TensorView>::Type == const TensorView
+//   MatchConstNess<TfLiteTensor, TensorView>::Type == TensorView
+//   MatchConstNess<TfLiteTensor, const TensorView>::Type == TensorView
+template <typename I, typename O>
+struct MatchConstNess {
+  using Type = std::conditional_t<std::is_const<I>::value, std::add_const_t<O>,
+                                  std::remove_const_t<O>>;
+};
+
+///////////////////////////// Implementation
+
+template <typename DType>
+DType &TensorView::AsScalar() {
+  DCHECK_EQ(shape_.size(), 0) << "Tensor is not a scalar";
+  return Data<DType>()[0];
+}
+
+template <typename DType>
+const DType &TensorView::AsScalar() const {
+  DCHECK_EQ(shape_.size(), 0) << "Tensor is not a scalar";
+  return Data<DType>().at(0);
+}
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TENSOR_VIEW_H_

tensorflow/lite/kernels/shim/test_op/BUILD

@@ -0,0 +1,161 @@
+# A simple op. for testing and demonstrating the OpKernel interface.
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_kernel_library")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//tensorflow/lite/kernels/shim:__subpackages__"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "simple_op",
+    hdrs = ["simple_op.h"],
+    deps = [
+        "//tensorflow/core/platform:tstring",
+        "//tensorflow/lite/kernels/shim:op_kernel",
+        "//tensorflow/lite/kernels/shim:shape",
+        "//tensorflow/lite/kernels/shim:status_macros",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_kernel_library(
+    name = "simple_tf_op",
+    srcs = ["simple_tf_op.cc"],
+    hdrs = ["simple_tf_op.h"],
+    deps = [
+        ":simple_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/lite/kernels/shim:tf_op_shim",
+    ],
+)
+
+tf_cc_test(
+    name = "simple_tf_op_test",
+    srcs = ["simple_tf_op_test.cc"],
+    tags = [
+        "nochromiumos_arm",
+        # Exclude from mobile builds
+        "tflite_not_portable",
+        "tflite_not_portable_android",
+    ],
+    deps = [
+        ":simple_tf_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:testlib",
+        "//tensorflow/core/kernels:ops_testutil",
+        "//tensorflow/core/platform:tstring",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ],
+)
+
+cc_library(
+    name = "simple_tflite_op",
+    srcs = ["simple_tflite_op.cc"],
+    hdrs = ["simple_tflite_op.h"],
+    deps = [
+        ":simple_op",
+        "//tensorflow/lite:mutable_op_resolver",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/shim:tflite_op_shim",
+    ],
+)
+
+cc_test(
+    name = "simple_tflite_op_test",
+    srcs = ["simple_tflite_op_test.cc"],
+    deps = [
+        ":simple_tflite_op",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "tmpl_op",
+    hdrs = ["tmpl_op.h"],
+    deps = [
+        "//tensorflow/lite/kernels/shim:op_kernel",
+        "//tensorflow/lite/kernels/shim:status_macros",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+tf_kernel_library(
+    name = "tmpl_tf_op",
+    srcs = ["tmpl_tf_op.cc"],
+    hdrs = ["tmpl_tf_op.h"],
+    deps = [
+        ":tmpl_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/lite/kernels/shim:tf_op_shim",
+    ],
+)
+
+tf_kernel_library(
+    name = "tmpl_tflite_op",
+    srcs = ["tmpl_tflite_op.cc"],
+    hdrs = ["tmpl_tflite_op.h"],
+    deps = [
+        ":tmpl_op",
+        "//tensorflow/lite:mutable_op_resolver",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/shim:op_kernel",
+        "//tensorflow/lite/kernels/shim:tflite_op_shim",
+        "//tensorflow/lite/kernels/shim:tflite_op_wrapper",
+    ],
+)
+
+tf_cc_test(
+    name = "tmpl_tf_op_test",
+    srcs = ["tmpl_tf_op_test.cc"],
+    tags = [
+        "nochromiumos_arm",
+        # Exclude from mobile builds
+        "tflite_not_portable",
+        "tflite_not_portable_android",
+    ],
+    deps = [
+        ":tmpl_tf_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/framework:fake_input",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/kernels:ops_testutil",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ],
+)
+
+tf_cc_test(
+    name = "tmpl_tflite_op_test",
+    srcs = ["tmpl_tflite_op_test.cc"],
+    tags = [
+        "nochromiumos_arm",
+        # Exclude from mobile builds
+        "tflite_not_portable",
+        "tflite_not_portable_android",
+    ],
+    deps = [
+        ":tmpl_tflite_op",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/kernels/shim/test_op/README.md

@@ -0,0 +1,19 @@
+This directory contains a fake operation in order to demonstrate and test the
+interfaces.
+
+First test op `SimpleOp` which is an op that test various attributes and input
+and output types. The other one is `TmplOp` which tests a templatized kernel.
+
+The contents:
+
+## `simple_op.h|cc`, `tmpl_op.h|cc`
+
+This is where the actual implementation of this op resides
+
+## `simple_tf_op.cc`, `tmpl_tf_op.cc`
+
+The TF op definition.
+
+## `simple_tflite_op.h|cc`
+
+The TFLite op definition.

tensorflow/lite/kernels/shim/test_op/simple_op.h

@@ -0,0 +1,208 @@
+/* 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_KERNELS_SHIM_TEST_OP_SIMPLE_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_OP_H_
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+
+namespace tflite {
+namespace shim {
+
+// A simple operation for demonstration and testing purposes.
+// See the kDoc member for documentation.
+
+template <Runtime Rt>
+class SimpleOp : public OpKernelShim<SimpleOp, Rt> {
+ protected:
+  enum Inputs { kInput0 = 0, kInput1 };
+  enum Outputs { kOutput0 = 0, kOutput1, kOutput2, kOutput3 };
+  int64_t output1_size_ = 0;
+  std::string output2_suffix_;
+  int64_t n_ = 0;
+  static constexpr int kOutput0Size = 5;
+  static constexpr char kOutput1SizeAttr[] = "output1_size";
+
+ public:
+  using typename OpKernelShim<SimpleOp, Rt>::InitContext;
+  using typename OpKernelShim<SimpleOp, Rt>::InvokeContext;
+  using typename OpKernelShim<SimpleOp, Rt>::ShapeInferenceContext;
+
+  SimpleOp() = default;
+  static constexpr char kOpName[] = "SimpleOperation";
+  static constexpr char kDoc[] = R"doc(
+Description:
+  Simple example op for testing and demonstration purposes.
+
+Attrs
+  output1_size: int - the size of the second output
+  output2_suffix: string - the string value to be appended to the end of out2
+  N: int - the number of dates for the second input and last output
+Inputs
+  in0: str, shape=[] - A scalar input
+  in1: int64, list<shape=?> - A list of tensors as input
+Outputs
+  out0: int, shape=[5] - first output
+  out1: float, shape=[?] - second output
+  out2: string, shape=[?] - third output
+  out3: int64, list<shape=?> - fourth output that is in1 but incremented.
+)doc";
+
+  static const char* OpName() { return kOpName; }
+  static const char* Doc() { return kDoc; }
+
+  // Attributes declaration (syntax: https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Attrs() {
+    return {absl::StrCat(kOutput1SizeAttr, ": int"), "output2_suffix: string",
+            "N: int >= 0"};
+  }
+  // Input tensors declaration (syntax:
+  // https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Inputs() {
+    return {"in0: string", "in1: N*int64"};
+  }
+  // Output tensors declaration (syntax:
+  // https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Outputs() {
+    return {"out0: int32", "out1: float", "out2: string", "out3: N*int64"};
+  }
+
+  // Initializes the op
+  absl::Status Init(InitContext* ctx) {
+    SH_RETURN_IF_ERROR(ctx->GetAttr(kOutput1SizeAttr, &output1_size_));
+    if (output1_size_ < 1 || output1_size_ > std::numeric_limits<int>::max()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat(kOutput1SizeAttr, " should be >= 1 and <= INT_MAX"));
+    }
+    SH_RETURN_IF_ERROR(ctx->GetAttr("N", &n_));
+    if (n_ < 0 || n_ > std::numeric_limits<int>::max()) {
+      return absl::InvalidArgumentError("N must be >= 0 and <= INT_MAX");
+    }
+    absl::string_view output2_suffix;
+    SH_RETURN_IF_ERROR(ctx->GetAttr("output2_suffix", &output2_suffix));
+    output2_suffix_ = output2_suffix;
+    return absl::OkStatus();
+  }
+
+  // Runs the operation
+  absl::Status Invoke(InvokeContext* ctx) {
+    using std::int32_t;
+    // read input
+    SH_ASSIGN_OR_RETURN(const auto input_t, ctx->GetInput(kInput0));
+    const tensorflow::tstring& input_str =
+        input_t->template AsScalar<tensorflow::tstring>();
+    // output0 whose size is static
+    SH_ASSIGN_OR_RETURN(auto output0_t,
+                        ctx->GetOutput(kOutput0, Shape({kOutput0Size})));
+    auto output0 = output0_t->template As<int32_t, 1>();
+    for (int i = 0; i < output0.Dim(0); ++i) output0(i) = i;
+    // output1 whose size is based on the attr
+    SH_ASSIGN_OR_RETURN(
+        auto output1_t,
+        ctx->GetOutput(kOutput1, Shape({static_cast<int>(output1_size_)})));
+    auto output1 = output1_t->template As<float, 1>();
+    for (int i = 0; i < output1.Dim(0); ++i) output1(i) = 0.5f * i;
+    // output2 whose size is based on input
+    if (input_str.length() >= std::numeric_limits<int>::max()) {
+      return absl::InvalidArgumentError("input_str length is too large.");
+    }
+    const int output2_size = input_str.length() + 1;
+    SH_ASSIGN_OR_RETURN(auto output2_t,
+                        ctx->GetOutput(kOutput2, Shape({output2_size})));
+    auto output2 = output2_t->template As<tensorflow::tstring, 1>();
+    for (int i = 0; i < output2.Dim(0) - 1; ++i) output2(i) = std::to_string(i);
+    if (output2.Dim(0) > 0) {
+      output2(output2.Dim(0) - 1) = output2_suffix_;
+    }
+    // output3 which is a list of length N
+    // The values in output3 are element wise equal to in1 + 1.
+    if (ctx->NumInputs() < kInput1 + n_) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "out of bounds: num_inputs=", ctx->NumInputs(), " N=", n_));
+    }
+    if (ctx->NumOutputs() < kOutput3 + n_) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "out of bounds: num_outputs=", ctx->NumOutputs(), " N=", n_));
+    }
+    for (int i = 0; i < n_; ++i) {
+      SH_ASSIGN_OR_RETURN(const auto input_t, ctx->GetInput(kInput1 + i));
+      Shape output_shape(input_t->Shape());
+      SH_ASSIGN_OR_RETURN(auto output_t,
+                          ctx->GetOutput(kOutput3 + i, output_shape));
+      const auto input_data = input_t->template Data<int64_t>();
+      auto& output_data = output_t->template Data<int64_t>();
+      const auto count = std::min(input_data.size(), output_data.size());
+      if (count > 0) {
+        std::transform(input_data.begin(), input_data.begin() + count,
+                       output_data.begin(), [](int64_t v) { return v + 1; });
+      }
+    }
+    return absl::OkStatus();
+  }
+
+  // Shape inference
+  static absl::Status ShapeInference(ShapeInferenceContext* ctx) {
+    int64_t n;
+    SH_RETURN_IF_ERROR(ctx->GetAttr("N", &n));
+    if (n + 1 != ctx->NumInputs()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "n + 1 != num_inputs: ", n + 1, " != ", ctx->NumInputs()));
+    }
+    if (n + 3 != ctx->NumOutputs()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "n + 3 != num_outputs: ", n + 3, " != ", ctx->NumOutputs()));
+    }
+    // output0
+    SH_RETURN_IF_ERROR(ctx->SetOutputShape(kOutput0, Shape({kOutput0Size})));
+    // output1
+    SH_RETURN_IF_ERROR(
+        ctx->SetOutputShape(kOutput1, Shape({Shape::kUnknownDim})));
+    // output2
+    const auto input_t_or = ctx->GetInputTensor(kInput0);
+    Shape output2_shape;
+    if (input_t_or.ok()) {
+      const auto& input_t = *input_t_or;
+      const auto& input_str = input_t->template AsScalar<tensorflow::tstring>();
+      if (input_str.length() >= std::numeric_limits<int>::max()) {
+        return absl::InvalidArgumentError("input_str length is too large.");
+      }
+      output2_shape = Shape({static_cast<int>(input_str.length() + 1)});
+    } else {
+      output2_shape = Shape({Shape::kUnknownDim});
+    }
+    SH_RETURN_IF_ERROR(ctx->SetOutputShape(kOutput2, output2_shape));
+    // output3
+    for (int i = kOutput3; i < ctx->NumOutputs(); ++i) {
+      SH_RETURN_IF_ERROR(ctx->SetOutputShape(i, Shape()));
+    }
+    return absl::OkStatus();
+  }
+};
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_OP_H_

tensorflow/lite/kernels/shim/test_op/simple_tf_op.cc

@@ -0,0 +1,31 @@
+/* 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/kernels/shim/test_op/simple_tf_op.h"
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/lite/kernels/shim/tf_op_shim.h"
+
+namespace tflite {
+namespace shim {
+
+REGISTER_TF_OP_SHIM(SimpleOpKernel);
+
+REGISTER_KERNEL_BUILDER(
+    Name(SimpleOpKernel::OpName()).Device(::tensorflow::DEVICE_CPU),
+    SimpleOpKernel);
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/simple_tf_op.h

@@ -0,0 +1,32 @@
+/* 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_KERNELS_SHIM_TEST_OP_SIMPLE_TF_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_TF_OP_H_
+
+#include "tensorflow/lite/kernels/shim/test_op/simple_op.h"
+#include "tensorflow/lite/kernels/shim/tf_op_shim.h"
+
+namespace tflite {
+namespace shim {
+
+class SimpleOpKernel : public TfOpKernel<SimpleOp> {
+ public:
+  using TfOpKernel::TfOpKernel;
+};
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_TF_OP_H_

tensorflow/lite/kernels/shim/test_op/simple_tf_op_test.cc

@@ -0,0 +1,97 @@
+/* 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 <cstdint>
+
+#include <gtest/gtest.h>
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/platform/tstring.h"
+
+namespace tflite {
+namespace shim {
+namespace {
+
+using ::tensorflow::DT_INT64;
+using ::tensorflow::DT_STRING;
+using ::tensorflow::FakeInput;
+using ::tensorflow::NodeDefBuilder;
+using ::tensorflow::TensorShape;
+using ::tensorflow::tstring;
+using ::tensorflow::test::AsTensor;
+using ::tensorflow::test::ExpectTensorEqual;
+
+class SimpleOpTfTest : public ::tensorflow::OpsTestBase {};
+
+TEST_F(SimpleOpTfTest, Output1Size_5_N_2) {
+  // Prepare graph.
+  TF_ASSERT_OK(NodeDefBuilder("simple_op", "SimpleOperation")
+                   .Attr("output1_size", 5)
+                   .Attr("output2_suffix", "foo")
+                   .Attr("N", 2)
+                   .Input(FakeInput(DT_STRING))
+                   .Input(FakeInput(2, DT_INT64))
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInputFromArray<tstring>(TensorShape({}), {"abc"});
+  AddInputFromArray<int64_t>(TensorShape({}), {123});
+  AddInputFromArray<int64_t>(TensorShape({2}), {456, 789});
+
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Validate the output.
+  ExpectTensorEqual<int>(*GetOutput(0),
+                         AsTensor<int>({0, 1, 2, 3, 4}, /*shape=*/{5}));
+  ExpectTensorEqual<float>(
+      *GetOutput(1), AsTensor<float>({0, 0.5, 1., 1.5, 2.}, /*shape=*/{5}));
+  ExpectTensorEqual<tstring>(
+      *GetOutput(2), AsTensor<tstring>({"0", "1", "2", "foo"}, /*shape=*/{4}));
+  ExpectTensorEqual<int64_t>(*GetOutput(3),
+                             AsTensor<int64_t>({124}, /*shape=*/{}));
+  ExpectTensorEqual<int64_t>(*GetOutput(4),
+                             AsTensor<int64_t>({457, 790}, /*shape=*/{2}));
+}
+
+TEST_F(SimpleOpTfTest, Output1Size_3_N_0) {
+  // Prepare graph.
+  TF_ASSERT_OK(NodeDefBuilder("simple_op", "SimpleOperation")
+                   .Attr("output1_size", 3)
+                   .Attr("output2_suffix", "foo")
+                   .Attr("N", 0)
+                   .Input(FakeInput(DT_STRING))
+                   .Input(FakeInput(0, DT_INT64))
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInputFromArray<tstring>(TensorShape({}), {"abcde"});
+
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Validate the output.
+  ExpectTensorEqual<int>(*GetOutput(0),
+                         AsTensor<int>({0, 1, 2, 3, 4}, /*shape=*/{5}));
+  ExpectTensorEqual<float>(*GetOutput(1),
+                           AsTensor<float>({0, 0.5, 1.}, /*shape=*/{3}));
+  ExpectTensorEqual<tstring>(
+      *GetOutput(2),
+      AsTensor<tstring>({"0", "1", "2", "3", "4", "foo"}, /*shape=*/{6}));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/simple_tflite_op.cc

@@ -0,0 +1,38 @@
+/* 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/kernels/shim/test_op/simple_tflite_op.h"
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/shim/test_op/simple_op.h"
+#include "tensorflow/lite/kernels/shim/tflite_op_shim.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+using OpKernel = ::tflite::shim::TfLiteOpKernel<tflite::shim::SimpleOp>;
+
+void AddSimpleOp(MutableOpResolver* resolver) { OpKernel::Add(resolver); }
+
+TfLiteRegistration* Register_SIMPLE_OP() {
+  return OpKernel::GetTfLiteRegistration();
+}
+
+const char* OpName_SIMPLE_OP() { return OpKernel::OpName(); }
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/simple_tflite_op.h

@@ -0,0 +1,37 @@
+/* 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_KERNELS_SHIM_TEST_OP_SIMPLE_TFLITE_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_TFLITE_OP_H_
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+// Add SimpleOp to the resolver
+void AddSimpleOp(MutableOpResolver* resolver);
+
+// Creates and returns the op kernel
+TfLiteRegistration* Register_SIMPLE_OP();
+
+// The name of the op
+const char* OpName_SIMPLE_OP();
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_SIMPLE_TFLITE_OP_H_

tensorflow/lite/kernels/shim/test_op/simple_tflite_op_test.cc

@@ -0,0 +1,140 @@
+/* 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/kernels/shim/test_op/simple_tflite_op.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace {
+
+class SimpleOpModel : public SingleOpModel {
+ public:
+  // Builds the op model and feeds in inputs, ready to invoke.
+  SimpleOpModel(const std::vector<uint8_t>& op_options,
+                const std::vector<tflite::TensorType>& input_types,
+                const std::vector<std::vector<int>>& input_shapes,
+                const std::string& input0,
+                const std::vector<std::vector<int64_t>>& input1,
+                const std::vector<tflite::TensorType>& output_types) {
+    // Define inputs.
+    std::vector<int> input_idx;
+    for (const auto input_type : input_types) {
+      input_idx.push_back(AddInput(input_type));
+    }
+    // Define outputs.
+    for (const auto output_type : output_types) {
+      output_idx_.push_back(AddOutput(output_type));
+    }
+    // Build the interpreter.
+    SetCustomOp(OpName_SIMPLE_OP(), op_options, Register_SIMPLE_OP);
+    BuildInterpreter(input_shapes);
+    // Populate inputs.
+    PopulateStringTensor(input_idx[0], {input0});
+    for (int i = 0; i < input1.size(); ++i) {
+      PopulateTensor(input_idx[1 + i], input1[i]);
+    }
+  }
+
+  template <typename T>
+  std::vector<T> GetOutput(const int i) {
+    return ExtractVector<T>(output_idx_[i]);
+  }
+
+  std::vector<int> GetOutputShape(const int i) {
+    return GetTensorShape(output_idx_[i]);
+  }
+
+ protected:
+  // Tensor indices
+  std::vector<int> output_idx_;
+};
+
+TEST(SimpleOpModel, OutputSize_5_N_2) {
+  // Test input
+  flexbuffers::Builder builder;
+  builder.Map([&]() {
+    builder.Int("output1_size", 5);
+    builder.String("output2_suffix", "foo");
+    builder.Int("N", 2);
+  });
+  builder.Finish();
+  std::vector<std::vector<int>> input_shapes = {{}, {}, {2}};
+  std::vector<tflite::TensorType> input_types = {tflite::TensorType_STRING,
+                                                 tflite::TensorType_INT64,
+                                                 tflite::TensorType_INT64};
+  std::vector<tflite::TensorType> output_types = {
+      tflite::TensorType_INT32, tflite::TensorType_FLOAT32,
+      tflite::TensorType_STRING, tflite::TensorType_INT64,
+      tflite::TensorType_INT64};
+  const std::string input0 = "abc";
+  const std::vector<std::vector<int64_t>> input1 = {{123}, {456, 789}};
+  // Run the op
+  SimpleOpModel m(/*op_options=*/builder.GetBuffer(), input_types, input_shapes,
+                  input0, input1, output_types);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  // Assertions
+  EXPECT_THAT(m.GetOutput<int>(0), testing::ElementsAre(0, 1, 2, 3, 4));
+  EXPECT_THAT(m.GetOutput<float>(1),
+              testing::ElementsAre(0, 0.5, 1.0, 1.5, 2.0));
+  EXPECT_THAT(m.GetOutput<std::string>(2),
+              testing::ElementsAre("0", "1", "2", "foo"));
+  EXPECT_THAT(m.GetOutput<int64_t>(3), testing::ElementsAre(124));
+  EXPECT_THAT(m.GetOutputShape(3), testing::ElementsAre());
+  EXPECT_THAT(m.GetOutput<int64_t>(4), testing::ElementsAre(457, 790));
+  EXPECT_THAT(m.GetOutputShape(4), testing::ElementsAre(2));
+}
+
+TEST(SimpleOpModel, OutputSize_3_N_0) {
+  // Test input
+  flexbuffers::Builder builder;
+  builder.Map([&]() {
+    builder.Int("output1_size", 3);
+    builder.String("output2_suffix", "foo");
+    builder.Int("N", 0);
+  });
+  builder.Finish();
+  std::vector<std::vector<int>> input_shapes = {{}};
+  std::vector<tflite::TensorType> input_types = {tflite::TensorType_STRING};
+  std::vector<tflite::TensorType> output_types = {tflite::TensorType_INT32,
+                                                  tflite::TensorType_FLOAT32,
+                                                  tflite::TensorType_STRING};
+  const std::string input0 = "abcde";
+  const std::vector<std::vector<int64_t>> input1;
+  // Run the op
+  SimpleOpModel m(/*op_options=*/builder.GetBuffer(), input_types, input_shapes,
+                  input0, input1, output_types);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  // Assertions
+  EXPECT_THAT(m.GetOutput<int>(0), testing::ElementsAre(0, 1, 2, 3, 4));
+  EXPECT_THAT(m.GetOutput<float>(1), testing::ElementsAre(0, 0.5, 1.0));
+  EXPECT_THAT(m.GetOutput<std::string>(2),
+              testing::ElementsAre("0", "1", "2", "3", "4", "foo"));
+}
+
+}  // namespace
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/tmpl_op.h

@@ -0,0 +1,104 @@
+/* Copyright 2022 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_KERNELS_SHIM_TEST_OP_TMPL_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_OP_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+
+namespace tflite {
+namespace shim {
+
+// A simple operation for demonstration and testing purposes.
+// See the kDoc member for documentation.
+
+template <Runtime Rt, typename AType, typename BType>
+class TmplOp : public OpKernelShim<TmplOp, Rt, AType, BType> {
+ protected:
+  enum Inputs { kInput0 = 0, kInput1 };
+  enum Outputs { kOutput0 = 0 };
+
+ public:
+  using typename OpKernelShim<TmplOp, Rt, AType, BType>::InitContext;
+  using typename OpKernelShim<TmplOp, Rt, AType, BType>::InvokeContext;
+  using typename OpKernelShim<TmplOp, Rt, AType, BType>::ShapeInferenceContext;
+
+  TmplOp() = default;
+  static constexpr char kOpName[] = "TemplatizedOperation";
+  static constexpr char kDoc[] = R"doc(
+Description:
+  Templatized op for testing and demonstration purposes.
+
+Attrs
+  AType: The type for input0
+  BType: The type for input1
+Inputs
+  in0: AType, shape=[] - A scalar input
+  in1: BType, shape=[] - A scalar input
+Outputs
+  out0: int, shape=[] - first output
+)doc";
+
+  static const char* OpName() { return kOpName; }
+  static const char* Doc() { return kDoc; }
+
+  // Attributes declaration (syntax: https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Attrs() {
+    return {"AType: {int32, float} = DT_INT32", "BType: type"};
+  }
+  // Input tensors declaration (syntax:
+  // https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Inputs() {
+    return {"in0: AType", "in1: BType"};
+  }
+  // Output tensors declaration (syntax:
+  // https://www.tensorflow.org/guide/create_op)
+  static std::vector<std::string> Outputs() { return {"out0: float"}; }
+
+  // Initializes the op
+  absl::Status Init(InitContext* ctx) { return absl::OkStatus(); }
+
+  // Shape inference
+  static absl::Status ShapeInference(ShapeInferenceContext* ctx) {
+    // outpu0
+    SH_RETURN_IF_ERROR(ctx->SetOutputShape(kOutput0, Shape({})));
+    return absl::OkStatus();
+  }
+
+  // Runs the operation
+  absl::Status Invoke(InvokeContext* ctx) {
+    using std::int32_t;
+    // input 0
+    SH_ASSIGN_OR_RETURN(const auto input0_t, ctx->GetInput(kInput0));
+    const auto in0 = input0_t->template AsScalar<AType>();
+    // input 1
+    SH_ASSIGN_OR_RETURN(const auto input1_t, ctx->GetInput(kInput1));
+    const auto in1 = input1_t->template AsScalar<BType>();
+    // output 0
+    SH_ASSIGN_OR_RETURN(auto output0_t, ctx->GetOutput(kOutput0, Shape({})));
+    auto& out0 = output0_t->template AsScalar<float>();
+    out0 = in0 + in1;
+    return absl::OkStatus();
+  }
+};
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_OP_H_

tensorflow/lite/kernels/shim/test_op/tmpl_tf_op.cc

@@ -0,0 +1,43 @@
+/* Copyright 2022 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/kernels/shim/test_op/tmpl_tf_op.h"
+
+#include <cstdint>
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/lite/kernels/shim/tf_op_shim.h"
+
+namespace tflite {
+namespace shim {
+
+using TmplOpKernelInstance = TmplOpKernel<float, int32_t>;
+
+REGISTER_TF_OP_SHIM(TmplOpKernelInstance);
+
+REGISTER_KERNEL_BUILDER(Name(TmplOpKernelInstance::OpName())
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<float>("AType")
+                            .TypeConstraint<int32_t>("BType"),
+                        TmplOpKernel<float, int32_t>);
+
+REGISTER_KERNEL_BUILDER(Name(TmplOpKernelInstance::OpName())
+                            .Device(::tensorflow::DEVICE_CPU)
+                            .TypeConstraint<int32_t>("AType")
+                            .TypeConstraint<int64_t>("BType"),
+                        TmplOpKernel<int32_t, int64_t>);
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/tmpl_tf_op.h

@@ -0,0 +1,33 @@
+/* Copyright 2022 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_KERNELS_SHIM_TEST_OP_TMPL_TF_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_TF_OP_H_
+
+#include "tensorflow/lite/kernels/shim/test_op/tmpl_op.h"
+#include "tensorflow/lite/kernels/shim/tf_op_shim.h"
+
+namespace tflite {
+namespace shim {
+
+template <typename AType, typename BType>
+class TmplOpKernel : public TfOpKernel<TmplOp, AType, BType> {
+ public:
+  using TfOpKernel<TmplOp, AType, BType>::TfOpKernel;
+};
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_TF_OP_H_

tensorflow/lite/kernels/shim/test_op/tmpl_tf_op_test.cc

@@ -0,0 +1,80 @@
+/* Copyright 2022 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 <cstdint>
+
+#include <gtest/gtest.h>
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+
+namespace tflite {
+namespace shim {
+namespace {
+
+using ::tensorflow::DT_FLOAT;
+using ::tensorflow::DT_INT32;
+using ::tensorflow::DT_INT64;
+using ::tensorflow::FakeInput;
+using ::tensorflow::NodeDefBuilder;
+using ::tensorflow::TensorShape;
+using ::tensorflow::test::AsTensor;
+using ::tensorflow::test::ExpectTensorEqual;
+
+class TmplOpTfTest : public ::tensorflow::OpsTestBase {};
+
+TEST_F(TmplOpTfTest, float_int32) {
+  // Prepare graph.
+  TF_ASSERT_OK(NodeDefBuilder("tmpl_op", "TemplatizedOperation")
+                   .Attr("AType", DT_FLOAT)
+                   .Attr("BType", DT_INT32)
+                   .Input(FakeInput(DT_FLOAT))
+                   .Input(FakeInput(DT_INT32))
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInputFromArray<float>(TensorShape({}), {10.5});
+  AddInputFromArray<int32_t>(TensorShape({}), {20});
+
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Validate the output.
+  ExpectTensorEqual<float>(*GetOutput(0),
+                           AsTensor<float>({30.5}, /*shape=*/{}));
+}
+
+TEST_F(TmplOpTfTest, int32_int64) {
+  // Prepare graph.
+  TF_ASSERT_OK(NodeDefBuilder("tmpl_op", "TemplatizedOperation")
+                   .Attr("AType", DT_INT32)
+                   .Attr("BType", DT_INT64)
+                   .Input(FakeInput(DT_INT32))
+                   .Input(FakeInput(DT_INT64))
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  AddInputFromArray<int32_t>(TensorShape({}), {10});
+  AddInputFromArray<int64_t>(TensorShape({}), {20});
+
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Validate the output.
+  ExpectTensorEqual<float>(*GetOutput(0), AsTensor<float>({30}, /*shape=*/{}));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/tmpl_tflite_op.cc

@@ -0,0 +1,53 @@
+/* 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/kernels/shim/test_op/tmpl_tflite_op.h"
+
+#include <cstdint>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/test_op/tmpl_op.h"
+#include "tensorflow/lite/kernels/shim/tflite_op_shim.h"
+#include "tensorflow/lite/kernels/shim/tflite_op_wrapper.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace {
+const char a_type[]("AType"), b_type[]("BType");
+}  // namespace
+
+using ::tflite::shim::op_wrapper::Attr;
+using ::tflite::shim::op_wrapper::AttrName;
+using ::tflite::shim::op_wrapper::OpWrapper;
+
+template <shim::Runtime Rt>
+using Op = OpWrapper<Rt, shim::TmplOp, Attr<AttrName<a_type>, int32_t, float>,
+                     Attr<AttrName<b_type>, int32_t, int64_t, bool>>;
+
+using OpKernel = ::tflite::shim::TfLiteOpKernel<Op>;
+
+void AddTmplOp(MutableOpResolver* resolver) { OpKernel::Add(resolver); }
+
+TfLiteRegistration* Register_TMPL_OP() {
+  return OpKernel::GetTfLiteRegistration();
+}
+
+const char* OpName_TMPL_OP() { return OpKernel::OpName(); }
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_op/tmpl_tflite_op.h

@@ -0,0 +1,38 @@
+/* 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_KERNELS_SHIM_TEST_OP_TMPL_TFLITE_OP_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_TFLITE_OP_H_
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+// Add TmplOp to the resolver
+void AddTmplOp(MutableOpResolver* resolver);
+
+// Creates and returns the op kernel
+TfLiteRegistration* Register_TMPL_OP();
+
+// The name of the op
+const char* OpName_TMPL_OP();
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_OP_TMPL_TFLITE_OP_H_

tensorflow/lite/kernels/shim/test_op/tmpl_tflite_op_test.cc

@@ -0,0 +1,144 @@
+/* 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/kernels/shim/test_op/tmpl_tflite_op.h"
+
+#include <cstdint>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace shim {
+namespace {
+
+template <typename AType, typename BType>
+class TmplOpModel : public SingleOpModel {
+ public:
+  // Builds the op model and feeds in inputs, ready to invoke.
+  TmplOpModel(const std::vector<uint8_t>& op_options,
+              const std::vector<tflite::TensorType>& input_types,
+              const std::vector<std::vector<int>>& input_shapes,
+              const std::vector<AType>& input0,
+              const std::vector<BType>& input1,
+              const std::vector<tflite::TensorType>& output_types) {
+    // Define inputs.
+    std::vector<int> input_idx;
+    for (const auto input_type : input_types) {
+      input_idx.push_back(AddInput(input_type));
+    }
+    // Define outputs.
+    for (const auto output_type : output_types) {
+      output_idx_.push_back(AddOutput(output_type));
+    }
+    // Build the interpreter.
+    SetCustomOp(ops::custom::OpName_TMPL_OP(), op_options,
+                ops::custom::Register_TMPL_OP);
+    BuildInterpreter(input_shapes);
+    // Populate inputs.
+    PopulateTensor(input_idx[0], input0);
+    PopulateTensor(input_idx[1], input1);
+  }
+
+  template <typename T>
+  std::vector<T> GetOutput(const int i) {
+    return ExtractVector<T>(output_idx_[i]);
+  }
+
+  std::vector<int> GetOutputShape(const int i) {
+    return GetTensorShape(output_idx_[i]);
+  }
+
+ protected:
+  // Tensor indices
+  std::vector<int> output_idx_;
+};
+
+TEST(TmplOpModel, float_int32) {
+  // Test input
+  flexbuffers::Builder builder;
+  builder.Map([&]() {
+    builder.Int("AType", kTfLiteFloat32);
+    builder.Int("BType", kTfLiteInt32);
+  });
+  builder.Finish();
+  std::vector<std::vector<int>> input_shapes = {{}, {}};
+  std::vector<tflite::TensorType> input_types = {tflite::TensorType_FLOAT32,
+                                                 tflite::TensorType_INT32};
+  std::vector<tflite::TensorType> output_types = {tflite::TensorType_FLOAT32};
+  const std::vector<float> input0 = {5.6f};
+  const std::vector<int32_t> input1 = {3};
+  // Run the op
+  TmplOpModel<float, int32_t> m(
+      /*op_options=*/builder.GetBuffer(), input_types, input_shapes, input0,
+      input1, output_types);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  // Assertions
+  EXPECT_THAT(m.GetOutput<float>(0), testing::ElementsAre(8.6f));
+}
+
+TEST(TmplOpModel, int32_int64) {
+  // Test input
+  flexbuffers::Builder builder;
+  builder.Map([&]() {
+    builder.Int("AType", kTfLiteInt32);
+    builder.Int("BType", kTfLiteInt64);
+  });
+  builder.Finish();
+  std::vector<std::vector<int>> input_shapes = {{}, {}};
+  std::vector<tflite::TensorType> input_types = {tflite::TensorType_INT32,
+                                                 tflite::TensorType_INT64};
+  std::vector<tflite::TensorType> output_types = {tflite::TensorType_FLOAT32};
+  const std::vector<int32_t> input0 = {12};
+  const std::vector<int64_t> input1 = {33l};
+  // Run the op
+  TmplOpModel<int32_t, int64_t> m(
+      /*op_options=*/builder.GetBuffer(), input_types, input_shapes, input0,
+      input1, output_types);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  // Assertions
+  EXPECT_THAT(m.GetOutput<float>(0), testing::ElementsAre(45.0f));
+}
+
+TEST(TmplOpModel, int32_bool) {
+  // Test input
+  flexbuffers::Builder builder;
+  builder.Map([&]() {
+    builder.Int("AType", kTfLiteInt32);
+    builder.Int("BType", kTfLiteBool);
+  });
+  builder.Finish();
+  std::vector<std::vector<int>> input_shapes = {{}, {}};
+  std::vector<tflite::TensorType> input_types = {tflite::TensorType_INT32,
+                                                 tflite::TensorType_BOOL};
+  std::vector<tflite::TensorType> output_types = {tflite::TensorType_FLOAT32};
+  const std::vector<int32_t> input0 = {12};
+  const std::vector<bool> input1 = {true};
+  // Run the op
+  TmplOpModel<int32_t, bool> m(
+      /*op_options=*/builder.GetBuffer(), input_types, input_shapes, input0,
+      input1, output_types);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  // Assertions
+  EXPECT_THAT(m.GetOutput<float>(0), testing::ElementsAre(13.0f));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_util.cc

@@ -0,0 +1,190 @@
+/* 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/kernels/shim/test_util.h"
+
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+
+using std::size_t;
+
+TfLiteTensor* UniqueTfLiteTensor::get() { return tensor_; }
+
+TfLiteTensor& UniqueTfLiteTensor::operator*() { return *tensor_; }
+
+TfLiteTensor* UniqueTfLiteTensor::operator->() { return tensor_; }
+
+const TfLiteTensor* UniqueTfLiteTensor::operator->() const { return tensor_; }
+
+void UniqueTfLiteTensor::reset(TfLiteTensor* tensor) { tensor_ = tensor; }
+
+UniqueTfLiteTensor::~UniqueTfLiteTensor() { TfLiteTensorFree(tensor_); }
+
+namespace {
+
+template <typename T>
+std::string TensorValueToString(const ::TfLiteTensor* tensor,
+                                const size_t idx) {
+  TFLITE_DCHECK_EQ(tensor->type, ::tflite::typeToTfLiteType<T>());
+  const T* val_array = reinterpret_cast<const T*>(tensor->data.raw);
+  return std::to_string(val_array[idx]);
+}
+
+template <>
+std::string TensorValueToString<bool>(const ::TfLiteTensor* tensor,
+                                      const size_t idx) {
+  TFLITE_DCHECK_EQ(tensor->type, ::tflite::typeToTfLiteType<bool>());
+  const bool* val_array = reinterpret_cast<const bool*>(tensor->data.raw);
+  return val_array[idx] ? "1" : "0";
+}
+
+template <typename FloatType>
+std::string TensorValueToStringFloat(const ::TfLiteTensor* tensor,
+                                     const size_t idx) {
+  TFLITE_DCHECK_EQ(tensor->type, ::tflite::typeToTfLiteType<FloatType>());
+  const FloatType* val_array =
+      reinterpret_cast<const FloatType*>(tensor->data.raw);
+  std::stringstream ss;
+  ss << val_array[idx];
+  return std::string(ss.str().data(), ss.str().length());
+}
+
+template <>
+std::string TensorValueToString<float>(const ::TfLiteTensor* tensor,
+                                       const size_t idx) {
+  return TensorValueToStringFloat<float>(tensor, idx);
+}
+
+template <>
+std::string TensorValueToString<double>(const ::TfLiteTensor* tensor,
+                                        const size_t idx) {
+  return TensorValueToStringFloat<double>(tensor, idx);
+}
+
+template <>
+std::string TensorValueToString<StringRef>(const ::TfLiteTensor* tensor,
+                                           const size_t idx) {
+  TFLITE_DCHECK_EQ(tensor->type, kTfLiteString);
+  const auto ref = ::tflite::GetString(tensor, idx);
+  return std::string(ref.str, ref.len);
+}
+
+std::string TfliteTensorDebugStringImpl(const ::TfLiteTensor* tensor,
+                                        const size_t axis,
+                                        const size_t max_values,
+                                        size_t* start_idx) {
+  const size_t dim_size = tensor->dims->data[axis];
+  if (axis == tensor->dims->size - 1) {
+    std::vector<std::string> ret_list;
+    ret_list.reserve(dim_size);
+    int idx = *start_idx;
+    for (int i = 0; i < dim_size && idx < max_values; ++i, ++idx) {
+      std::string val_str;
+      switch (tensor->type) {
+        case kTfLiteBool: {
+          val_str = TensorValueToString<bool>(tensor, idx);
+          break;
+        }
+        case kTfLiteUInt8: {
+          val_str = TensorValueToString<uint8_t>(tensor, idx);
+          break;
+        }
+        case kTfLiteInt8: {
+          val_str = TensorValueToString<int8_t>(tensor, idx);
+          break;
+        }
+        case kTfLiteInt16: {
+          val_str = TensorValueToString<int16_t>(tensor, idx);
+          break;
+        }
+        case kTfLiteInt32: {
+          val_str = TensorValueToString<int32_t>(tensor, idx);
+          break;
+        }
+        case kTfLiteInt64: {
+          val_str = TensorValueToString<int64_t>(tensor, idx);
+          break;
+        }
+        case kTfLiteString: {
+          val_str = TensorValueToString<StringRef>(tensor, idx);
+          break;
+        }
+        case kTfLiteFloat32: {
+          val_str = TensorValueToString<float>(tensor, idx);
+          break;
+        }
+        case kTfLiteFloat64: {
+          val_str = TensorValueToString<double>(tensor, idx);
+          break;
+        }
+        default: {
+          val_str = "unsupported_type";
+        }
+      }
+      ret_list.push_back(val_str);
+    }
+    *start_idx = idx;
+    if (idx == max_values && ret_list.size() < dim_size) {
+      ret_list.push_back("...");
+    }
+    return absl::StrCat("[", absl::StrJoin(ret_list, ", "), "]");
+  } else {
+    std::vector<std::string> ret_list;
+    ret_list.reserve(dim_size);
+    for (int i = 0; i < dim_size && *start_idx < max_values; ++i) {
+      ret_list.push_back(
+          TfliteTensorDebugStringImpl(tensor, axis + 1, max_values, start_idx));
+    }
+    return absl::StrCat("[", absl::StrJoin(ret_list, ", "), "]");
+  }
+}
+
+}  // namespace
+
+std::string TfliteTensorDebugString(const ::TfLiteTensor* tensor,
+                                    const size_t max_values) {
+  if (tensor->dims->size == 0) return "";
+  size_t start_idx = 0;
+  return TfliteTensorDebugStringImpl(tensor, 0, max_values, &start_idx);
+}
+
+size_t NumTotalFromShape(const std::initializer_list<int>& shape) {
+  size_t num_total;
+  if (shape.size() > 0)
+    num_total = 1;
+  else
+    num_total = 0;
+  for (const int dim : shape) num_total *= dim;
+  return num_total;
+}
+
+template <>
+void PopulateTfLiteTensorValue<std::string>(
+    const std::initializer_list<std::string> values, TfLiteTensor* tensor) {
+  tflite::DynamicBuffer buf;
+  for (const std::string& s : values) {
+    buf.AddString(s.data(), s.length());
+  }
+  buf.WriteToTensor(tensor, /*new_shape=*/nullptr);
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/shim/test_util.h

@@ -0,0 +1,111 @@
+/* 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_KERNELS_SHIM_TEST_UTIL_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TEST_UTIL_H_
+
+#include <initializer_list>
+#include <string>
+
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/type_to_tflitetype.h"
+
+namespace tflite {
+
+// A wrapper around TfLiteTensor which frees it in its dtor.
+class UniqueTfLiteTensor {
+ public:
+  explicit UniqueTfLiteTensor(TfLiteTensor* tensor) : tensor_(tensor) {}
+
+  UniqueTfLiteTensor() = default;
+
+  // Returns the underlying pointer
+
+  TfLiteTensor* get();
+
+  TfLiteTensor& operator*();
+
+  TfLiteTensor* operator->();
+
+  const TfLiteTensor* get() const;
+
+  const TfLiteTensor& operator*() const;
+
+  const TfLiteTensor* operator->() const;
+
+  // Resets the underlying pointer
+  void reset(TfLiteTensor* tensor);
+
+  // Deallocates the tensor as well
+  ~UniqueTfLiteTensor();
+
+ private:
+  TfLiteTensor* tensor_ = nullptr;
+};
+
+// Prints a debug string for the given tensor.
+std::string TfliteTensorDebugString(const ::TfLiteTensor* tensor,
+                                    const std::size_t max_values = 30);
+
+// Calculate the total number of elements given the shape.
+std::size_t NumTotalFromShape(const std::initializer_list<int>& shape);
+
+template <typename T>
+void ReallocDynamicTensor(const std::initializer_list<int> shape,
+                          TfLiteTensor* tensor) {
+  TfLiteTensorFree(tensor);
+  tensor->allocation_type = kTfLiteDynamic;
+  tensor->type = typeToTfLiteType<T>();
+  // Populate Shape
+  TfLiteIntArray* shape_arr = TfLiteIntArrayCreate(shape.size());
+  int i = 0;
+  const std::size_t num_total = NumTotalFromShape(shape);
+  for (const int dim : shape) shape_arr->data[i++] = dim;
+  tensor->dims = shape_arr;
+  if (tensor->type != kTfLiteString) {
+    TfLiteTensorRealloc(num_total * sizeof(T), tensor);
+  }
+}
+
+// Populates a tensor with the given values
+template <typename T>
+void PopulateTfLiteTensorValue(const std::initializer_list<T> values,
+                               TfLiteTensor* tensor) {
+  T* buffer = reinterpret_cast<T*>(tensor->data.raw);
+  int i = 0;
+  for (const auto v : values) {
+    buffer[i++] = v;
+  }
+}
+
+template <>
+void PopulateTfLiteTensorValue<std::string>(
+    const std::initializer_list<std::string> values, TfLiteTensor* tensor);
+
+template <typename T>
+void PopulateTfLiteTensor(const std::initializer_list<T> values,
+                          const std::initializer_list<int> shape,
+                          TfLiteTensor* tensor) {
+  const std::size_t num_total = NumTotalFromShape(shape);
+  CHECK_EQ(num_total, values.size());
+  // Populate Shape
+  ReallocDynamicTensor<T>(shape, tensor);
+  // Value allocation
+  PopulateTfLiteTensorValue<T>(values, tensor);
+}
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TEST_UTIL_H_

tensorflow/lite/kernels/shim/test_util_test.cc

@@ -0,0 +1,60 @@
+/* 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/kernels/shim/test_util.h"
+
+#include <string>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+
+namespace tflite {
+namespace {
+
+TEST(TfliteTensorDebugString, Basic) {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(3);
+  interpreter.AllocateTensors();
+
+  auto t_int32 = UniqueTfLiteTensor(interpreter.tensor(0));
+  PopulateTfLiteTensor<int32_t>({1, 2, 3, 4, 5}, {5}, t_int32.get());
+  EXPECT_EQ("[1, 2, 3, 4, 5]", TfliteTensorDebugString(t_int32.get()));
+
+  auto t_int64 = UniqueTfLiteTensor(interpreter.tensor(1));
+  PopulateTfLiteTensor<int32_t>({1, 2, 3, 4}, {2, 2}, t_int64.get());
+  EXPECT_EQ("[[1, 2], [3, 4]]", TfliteTensorDebugString(t_int64.get()));
+
+  auto t_str = UniqueTfLiteTensor(interpreter.tensor(2));
+  PopulateTfLiteTensor<std::string>({"ab", "cde", "f"}, {1, 3}, t_str.get());
+  EXPECT_EQ("[[ab, cde, f]]", TfliteTensorDebugString(t_str.get()));
+}
+
+TEST(TfliteTensorDebugString, MaxVal) {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(2);
+  interpreter.AllocateTensors();
+
+  auto t_int32 = UniqueTfLiteTensor(interpreter.tensor(0));
+  PopulateTfLiteTensor<int32_t>({1, 2, 3, 4}, {4}, t_int32.get());
+  EXPECT_EQ("[1, 2, 3, 4]",
+            TfliteTensorDebugString(t_int32.get(), /*max_values=*/4));
+
+  t_int32 = UniqueTfLiteTensor(interpreter.tensor(0));
+  PopulateTfLiteTensor<int32_t>({1, 2, 3, 4, 5}, {5}, t_int32.get());
+  EXPECT_EQ("[1, 2, 3, 4, ...]",
+            TfliteTensorDebugString(t_int32.get(), /*max_values=*/4));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/shim/tf_op_shim.cc

@@ -0,0 +1,179 @@
+/* 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/kernels/shim/tf_op_shim.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+#include "tensorflow/lite/kernels/shim/tf_tensor_view.h"
+
+namespace tflite {
+namespace shim {
+
+namespace {
+// Converts a TF AttrValue into a TF Shim AttrValue
+absl::StatusOr<AttrValue> TfAttrValueToShimAttrValue(
+    const ::tensorflow::AttrValue& attr_value) {
+  AttrValue ret;
+  switch (attr_value.value_case()) {
+    case ::tensorflow::AttrValue::kB: {
+      ret = attr_value.b();
+      break;
+    }
+    case ::tensorflow::AttrValue::kI: {
+      ret = attr_value.i();
+      break;
+    }
+    case ::tensorflow::AttrValue::kF: {
+      ret = attr_value.f();
+      break;
+    }
+    case ::tensorflow::AttrValue::kS: {
+      ret = attr_value.s();
+      break;
+    }
+    default: {
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Unsupported attribute type: ", attr_value.DebugString()));
+    }
+  }
+  return ret;
+}
+}  // namespace
+
+TfInitContext::TfInitContext(const ::tensorflow::OpKernelConstruction* context)
+    : context_(context) {}
+
+absl::StatusOr<AttrValue> TfInitContext::GetAttr(
+    const std::string& attr_name) const {
+  if (!context_->HasAttr(attr_name))
+    return absl::InvalidArgumentError(
+        absl::StrCat("Non-existent attribute: ", attr_name, "\nop def:\n",
+                     context_->def().DebugString()));
+  const auto& attr_value = context_->def().attr().at(attr_name);
+  return TfAttrValueToShimAttrValue(attr_value);
+}
+
+TfInvokeContext::TfInvokeContext(::tensorflow::OpKernelContext* context)
+    : context_(context) {}
+
+ConstTensorViewOr TfInvokeContext::GetInput(const int idx) const {
+  if (idx >= context_->num_inputs()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expected idx < num_inputs. idx: ", idx,
+                     " num_inputs: ", context_->num_inputs()));
+  }
+  const auto tf_tensor = context_->input(idx);
+  SH_ASSIGN_OR_RETURN(const TfTensorView& tensor_view,
+                      TensorView::New(&tf_tensor));
+  return std::make_unique<const TfTensorView>(tensor_view);
+}
+
+TensorViewOr TfInvokeContext::GetOutput(const int idx,
+                                        const Shape& shape) const {
+  tensorflow::Tensor* output_t = nullptr;
+  if (!shape.has_value())
+    return absl::InvalidArgumentError("Output shape needs to be specified.");
+  std::vector<int64_t> shape_64(shape->size());
+  for (int i = 0; i < shape->size(); ++i) shape_64[i] = (*shape)[i];
+  auto status = context_->allocate_output(
+      idx, ::tensorflow::TensorShape(shape_64), &output_t);
+  if (!status.ok()) return status;
+  SH_ASSIGN_OR_RETURN(const TfTensorView& tensor_view,
+                      TensorView::New(output_t));
+  return std::make_unique<TfTensorView>(std::move(tensor_view));
+}
+
+int TfInvokeContext::NumInputs() const { return context_->num_inputs(); }
+
+int TfInvokeContext::NumOutputs() const { return context_->num_outputs(); }
+
+TfShapeInferenceContext::TfShapeInferenceContext(
+    ::tensorflow::shape_inference::InferenceContext* context)
+    : context_(context) {}
+
+ShapeOr TfShapeInferenceContext::GetInputShape(const int idx) const {
+  std::vector<int> ret;
+  const auto& shape = context_->input(idx);
+  if (!context_->RankKnown(shape)) return Shape();
+  ret.resize(context_->Rank(shape));
+  for (int i = 0; i < ret.size(); ++i)
+    ret[i] = context_->Value(context_->Dim(shape, i));
+  return Shape(ret);
+}
+
+absl::Status TfShapeInferenceContext::SetOutputShape(const int idx,
+                                                     const Shape& shape) {
+  tensorflow::shape_inference::ShapeHandle output_shape;
+  if (shape.has_value()) {
+    std::vector<::tensorflow::shape_inference::DimensionHandle> tf_shape;
+    tf_shape.reserve(shape.value().size());
+    for (const auto dim : shape.value())
+      tf_shape.emplace_back(context_->MakeDim(dim));
+    output_shape = context_->MakeShape(tf_shape);
+  } else {
+    output_shape = context_->UnknownShape();
+  }
+  context_->set_output(idx, output_shape);
+  return absl::OkStatus();
+}
+
+ConstTensorViewOr TfShapeInferenceContext::GetInputTensor(const int idx) const {
+  const auto* tf_tensor = context_->input_tensor(idx);
+  if (tf_tensor == nullptr) {
+    return absl::UnavailableError(
+        absl::StrCat("Tensor is not available. idx: ", idx));
+  }
+  SH_ASSIGN_OR_RETURN(const TfTensorView& tensor_view,
+                      TensorView::New(tf_tensor));
+  return std::make_unique<const TfTensorView>(tensor_view);
+}
+
+absl::StatusOr<AttrValue> TfShapeInferenceContext::GetAttr(
+    const std::string& attr_name) const {
+  const auto* tf_attr_value = context_->GetAttr(attr_name);
+  if (tf_attr_value == nullptr)
+    return absl::InvalidArgumentError(
+        absl::StrCat("Non-existent attribute: ", attr_name));
+  return TfAttrValueToShimAttrValue(*tf_attr_value);
+}
+
+int TfShapeInferenceContext::NumInputs() const {
+  return context_->num_inputs();
+}
+
+int TfShapeInferenceContext::NumOutputs() const {
+  return context_->num_outputs();
+}
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/tf_op_shim.h

@@ -0,0 +1,166 @@
+/* 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_KERNELS_SHIM_TF_OP_SHIM_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TF_OP_SHIM_H_
+
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tsl/platform/macros.h"
+
+// This file contains the TF adapter. That is, it takes a `OpKernelShim`
+// class and provides a TF kernel out of it.
+
+namespace tflite {
+namespace shim {
+
+// TF implementation of the methods during an op kernel initialization
+class TfInitContext : public InitContext<TfInitContext> {
+ public:
+  explicit TfInitContext(const ::tensorflow::OpKernelConstruction* context);
+  // Read a given attribute
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const;
+
+ private:
+  const ::tensorflow::OpKernelConstruction* context_;
+};
+
+// TF implementation of the methods during an op kernel invocation
+class TfInvokeContext : public InvokeContext<TfInvokeContext> {
+ public:
+  explicit TfInvokeContext(::tensorflow::OpKernelContext* context);
+  // Read an input tensor
+  ConstTensorViewOr GetInput(int idx) const;
+  // Get a mutable output tensor
+  TensorViewOr GetOutput(int idx, const Shape& shape) const;
+  // Number of input tensors
+  int NumInputs() const;
+  // Number of output tensors
+  int NumOutputs() const;
+
+ private:
+  ::tensorflow::OpKernelContext* context_;
+};
+
+// TF implementation of the methods during shape inference
+class TfShapeInferenceContext
+    : public ShapeInferenceContext<TfShapeInferenceContext> {
+ public:
+  explicit TfShapeInferenceContext(
+      ::tensorflow::shape_inference::InferenceContext* context);
+  // Read an input tensor shape
+  ShapeOr GetInputShape(int idx) const;
+  // Set an output tensor shape
+  absl::Status SetOutputShape(int idx, const Shape& shape);
+  // Read an input tensor during shape inference
+  ConstTensorViewOr GetInputTensor(int idx) const;
+  // Read a given attribute
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const;
+  // Number of input tensors
+  int NumInputs() const;
+  // Number of output tensors
+  int NumOutputs() const;
+
+ private:
+  ::tensorflow::shape_inference::InferenceContext* context_;
+};
+
+// The adaptor between an op implementation (OpKernelShim subclass) and TF
+// runtime
+template <template <Runtime, typename...> typename Impl, typename... Ts>
+class TfOpKernel : public ::tensorflow::OpKernel {
+ public:
+  using ImplType = Impl<Runtime::kTf, Ts...>;
+
+  explicit TfOpKernel(::tensorflow::OpKernelConstruction* c)
+      : OpKernel(c), impl_(std::make_unique<ImplType>()) {
+    TfInitContext ctx(c);
+    c->SetStatus(impl_->Init(&ctx));
+  }
+
+  // The main computation of the op
+  void Compute(::tensorflow::OpKernelContext* c) override {
+    TfInvokeContext ctx(c);
+    OP_REQUIRES_OK(c, impl_->Invoke(&ctx));
+  }
+
+  // Shape inference for the op.
+  static absl::Status ShapeInference(
+      ::tensorflow::shape_inference::InferenceContext* c) {
+    TfShapeInferenceContext ctx(c);
+    return ImplType::ShapeInference(&ctx);
+  }
+
+  // The operation name
+  static const char* OpName() { return ImplType::OpName(); }
+
+ protected:
+  std::unique_ptr<OpKernelShim<Impl, Runtime::kTf, Ts...>> impl_;
+};
+
+static_assert(::tensorflow::shape_inference::InferenceContext::kUnknownDim ==
+                  Shape::kUnknownDim,
+              "The values must match.");
+static_assert(::tensorflow::shape_inference::InferenceContext::kUnknownRank ==
+                  Shape::kUnknownRank,
+              "The values must match.");
+
+// Builds the OpDef to register the op with the TF runtime
+template <typename Kernel>
+::tensorflow::register_op::OpDefBuilderWrapper CreateOpDefBuilderWrapper() {
+  auto ret = ::tensorflow::register_op::OpDefBuilderWrapper(
+      Kernel::ImplType::OpName());
+  for (const auto& input : Kernel::ImplType::Inputs()) ret = ret.Input(input);
+  for (const auto& output : Kernel::ImplType::Outputs())
+    ret = ret.Output(output);
+  for (const auto& attr : Kernel::ImplType::Attrs()) ret = ret.Attr(attr);
+  ret.SetShapeFn(Kernel::ShapeInference).Doc(Kernel::ImplType::kDoc);
+  return ret;
+}
+
+template <>
+struct ContextTypeForRuntime<Runtime::kTf> {
+  using Init = TfInitContext;
+  using Invoke = TfInvokeContext;
+  using ShapeInference = TfShapeInferenceContext;
+};
+
+// Macros for defining an op. These are taken from op.h because they need to be
+// slightly modified here.
+#define REGISTER_OP_SHIM_IMPL(ctr, op_kernel_cls)                            \
+  static ::tensorflow::InitOnStartupMarker const register_op##ctr            \
+      TF_ATTRIBUTE_UNUSED =                                                  \
+          TF_INIT_ON_STARTUP_IF(SHOULD_REGISTER_OP(op_kernel_cls::OpName())) \
+          << ::tflite::shim::CreateOpDefBuilderWrapper<op_kernel_cls>()
+
+#define REGISTER_TF_OP_SHIM(...) \
+  TF_ATTRIBUTE_ANNOTATE("tf:op") \
+  TF_NEW_ID_FOR_INIT(REGISTER_OP_SHIM_IMPL, __VA_ARGS__)
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TF_OP_SHIM_H_

tensorflow/lite/kernels/shim/tf_tensor_view.cc

@@ -0,0 +1,99 @@
+/* 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/kernels/shim/tf_tensor_view.h"
+
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+// Creates a case statement for the switch() clause given the dtype
+#define CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(TF_DTYPE, CPP_DTYPE) \
+  case TF_DTYPE: {                                          \
+    using DType = typename CPP_DTYPE;                       \
+    return TfTensorView(wrapped_tensor, DType());           \
+  }
+
+#define CASE_FOR_DTYPE(TF_DTYPE)           \
+  CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(TF_DTYPE, \
+                                 ::tensorflow::EnumToDataType<TF_DTYPE>::Type)
+
+namespace tflite {
+namespace shim {
+
+// ctors
+
+TfTensorView::TfTensorView(TfTensorView &&o) noexcept
+    : TensorView(std::move(o)), shape_data_(std::move(o.shape_data_)) {
+  shape_ = absl::Span<int>(shape_data_);
+}
+
+TfTensorView::TfTensorView(const TfTensorView &o)
+    : TensorView(o), shape_data_(o.shape_data_) {
+  shape_ = absl::Span<int>(shape_data_);
+}
+
+TfTensorView &TfTensorView::operator=(TfTensorView &&o) noexcept {
+  shape_data_ = std::move(o.shape_data_);
+  TensorView::operator=(std::move(o));
+  shape_ = absl::Span<int>(shape_data_);
+  return *this;
+}
+
+TfTensorView &TfTensorView::operator=(const TfTensorView &o) {
+  if (&o == this) return *this;
+  TensorView::operator=(o);
+  shape_data_ = o.shape_data_;
+  shape_ = absl::Span<int>(shape_data_);
+  return *this;
+}
+
+template <typename TfTensorType>
+absl::StatusOr<typename MatchConstNess<TfTensorType, TfTensorView>::Type>
+TfTensorViewTemplatizedNew(TfTensorType *wrapped_tensor) {
+  switch (wrapped_tensor->dtype()) {
+    CASE_FOR_DTYPE(::tensorflow::DT_BOOL);
+    CASE_FOR_DTYPE(::tensorflow::DT_UINT8);
+    CASE_FOR_DTYPE(::tensorflow::DT_UINT64);
+    CASE_FOR_DTYPE(::tensorflow::DT_INT8);
+    CASE_FOR_DTYPE(::tensorflow::DT_INT16);
+    CASE_FOR_DTYPE(::tensorflow::DT_INT32);
+    // Map DT_INT64 to int64_t instead of int64 to have a single int64 datatype.
+    CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(::tensorflow::DT_INT64, std::int64_t);
+    CASE_FOR_DTYPE(::tensorflow::DT_FLOAT);
+    CASE_FOR_DTYPE(::tensorflow::DT_DOUBLE);
+    CASE_FOR_DTYPE(::tensorflow::DT_STRING);
+    default: {
+      return absl::UnimplementedError(
+          absl::StrCat("Unsupported data type: ", wrapped_tensor->dtype()));
+    }
+  }
+}
+
+template <>
+absl::StatusOr<TfTensorView> TensorView::New<::tensorflow::Tensor>(
+    ::tensorflow::Tensor *wrapped_tensor) {
+  return TfTensorViewTemplatizedNew(wrapped_tensor);
+}
+
+template <>
+absl::StatusOr<const TfTensorView> TensorView::New<const ::tensorflow::Tensor>(
+    const ::tensorflow::Tensor *wrapped_tensor) {
+  return TfTensorViewTemplatizedNew(wrapped_tensor);
+}
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/tf_tensor_view.h

@@ -0,0 +1,99 @@
+/* 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_KERNELS_SHIM_TF_TENSOR_VIEW_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TF_TENSOR_VIEW_H_
+
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+
+namespace tflite {
+namespace shim {
+
+// A view over TF Tensor without taking ownership. It can be either mutable or
+// immutable.
+class TfTensorView : public TensorView {
+ public:
+  // Move constructor
+  TfTensorView(TfTensorView &&o) noexcept;
+  // Copy constructor
+  TfTensorView(const TfTensorView &o);
+  // Move assignment operator
+  TfTensorView &operator=(TfTensorView &&o) noexcept;
+  // Copy assignment operator
+  TfTensorView &operator=(const TfTensorView &);
+
+ protected:
+  // Templated constructor. Since it's not possible to specify the template
+  // argument directly we place a dummy argument of that type so compiler can
+  // deduce the right template parameter
+  template <typename DType>
+  TfTensorView(const ::tensorflow::Tensor *wrapped_tensor, const DType &dtype);
+
+  // Let the factory implementation use private constructors
+  template <typename TfTensorType>
+  friend absl::StatusOr<
+      typename MatchConstNess<TfTensorType, TfTensorView>::Type>
+  TfTensorViewTemplatizedNew(TfTensorType *wrapped_tensor);
+
+  // Stores the shape read from the TensorShape object
+  std::vector<int> shape_data_;
+};
+
+// Map ::tensorflow::Tensor -> TfTensorView
+template <>
+struct TensorViewSubType<::tensorflow::Tensor> {
+  using Type = TfTensorView;
+};
+
+// Map const ::tensorflow::Tensor -> const TfTensorView
+template <>
+struct TensorViewSubType<const ::tensorflow::Tensor> {
+  using Type = const TfTensorView;
+};
+
+// Specialization of New() factory
+template <>
+absl::StatusOr<TfTensorView> TensorView::New<::tensorflow::Tensor>(
+    ::tensorflow::Tensor *wrapped_tensor);
+
+// Specialization of New() factory
+template <>
+absl::StatusOr<const TfTensorView> TensorView::New<const ::tensorflow::Tensor>(
+    const ::tensorflow::Tensor *wrapped_tensor);
+
+/////////////////////// Implementation
+///////////////////////
+
+// Templated ctor
+template <typename DType>
+TfTensorView::TfTensorView(const ::tensorflow::Tensor *wrapped_tensor,
+                           const DType &dtype)
+    : TensorView({}, wrapped_tensor->data(), wrapped_tensor->TotalBytes(),
+                 dtype) {
+  shape_data_.resize(wrapped_tensor->shape().dims());
+  for (int dim = 0; dim < wrapped_tensor->shape().dims(); ++dim) {
+    shape_data_[dim] = wrapped_tensor->shape().dim_size(dim);
+  }
+  shape_ = absl::Span<int>(shape_data_);
+}
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TF_TENSOR_VIEW_H_

tensorflow/lite/kernels/shim/tf_tensor_view_test.cc

@@ -0,0 +1,217 @@
+/* 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/kernels/shim/tf_tensor_view.h"
+
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/tstring.h"
+
+namespace tflite {
+namespace shim {
+namespace {
+
+using ::tensorflow::protobuf::TextFormat;
+
+TEST(TfTensorView, Bool) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_BOOL
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        bool_val: [ false, false, false, false, false, false ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+
+  // Test move assignment
+  auto t_premove_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_premove_or.ok()) << t_premove_or.status();
+  auto t = std::move(t_premove_or.value());
+
+  auto tensor_data_as_vector = t.Data<bool>();
+  for (int i = 0; i < 3 * 2; ++i) tensor_data_as_vector[i] = i % 5 == 0;
+
+  ASSERT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq("[[1 0]\n [0 0]\n [0 1]]"));
+}
+
+TEST(TfTensorView, Int32) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_INT32
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        int_val: [ 0, 0, 0, 0, 0, 0 ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+
+  // Test move assignment
+  auto t_premove_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_premove_or.ok()) << t_premove_or.status();
+  auto t = std::move(t_premove_or.value());
+
+  auto tensor_data_as_vector = t.Data<int32_t>();
+  for (int i = 0; i < 3 * 2; ++i) tensor_data_as_vector[i] = i;
+
+  ASSERT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq("[[0 1]\n [2 3]\n [4 5]]"));
+}
+
+TEST(TfTensorView, Int64) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_INT64
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        int_val: [ 0, 0, 0, 0, 0, 0 ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+  auto t_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_or.ok()) << t_or.status();
+  auto& t = t_or.value();
+
+  auto tensor_data_as_vector = t.Data<int64_t>();
+  for (int i = 0; i < 3 * 2; ++i) tensor_data_as_vector[i] = i;
+
+  ASSERT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq("[[0 1]\n [2 3]\n [4 5]]"));
+}
+
+TEST(TfTensorView, Float) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_FLOAT
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        float_val: [ 0, 0, 0, 0, 0, 0 ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+  auto t_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_or.ok()) << t_or.status();
+  auto& t = t_or.value();
+
+  auto tensor_data_as_vector = t.Data<float>();
+  for (int i = 0; i < 3 * 2; ++i)
+    tensor_data_as_vector[i] = static_cast<float>(i) / 2.0;
+
+  ASSERT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq("[[0 0.5]\n [1 1.5]\n [2 2.5]]"));
+}
+
+TEST(TfTensorView, Double) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_DOUBLE
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        double_val: [ 0, 0, 0, 0, 0, 0 ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+  auto t_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_or.ok()) << t_or.status();
+  auto& t = t_or.value();
+
+  auto tensor_data_as_vector = t.Data<double>();
+  for (int i = 0; i < 3 * 2; ++i)
+    tensor_data_as_vector[i] = static_cast<double>(i) / 2.0;
+
+  ASSERT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq("[[0 0.5]\n [1 1.5]\n [2 2.5]]"));
+}
+
+TEST(TfTensorView, Str) {
+  ::tensorflow::TensorProto tf_tensor_pb;
+  ASSERT_TRUE(TextFormat::ParseFromString(
+      R"pb(
+        dtype: DT_STRING
+        tensor_shape {
+          dim:
+          [ { size: 3 }
+            , { size: 2 }]
+        }
+        string_val: [ "", "", "", "", "", "" ]
+      )pb",
+      &tf_tensor_pb));
+  ::tensorflow::Tensor tf_tensor;
+  ASSERT_TRUE(tf_tensor.FromProto(tf_tensor_pb));
+  auto t_or = TensorView::New(&tf_tensor);
+  ASSERT_TRUE(t_or.ok()) << t_or.status();
+  auto& t = t_or.value();
+
+  auto tensor_data_as_vector = t.Data<::tensorflow::tstring>();
+  tensor_data_as_vector[0] = "a";
+  tensor_data_as_vector[1] = "bc";
+  tensor_data_as_vector[2] = "def";
+  tensor_data_as_vector[3] = "g";
+  tensor_data_as_vector[4] = "hi";
+  tensor_data_as_vector[5] = "";
+
+  EXPECT_THAT(t.Data<::tensorflow::tstring>(),
+              ::testing::ElementsAre("a", "bc", "def", "g", "hi", ""));
+
+  const auto& const_tf_tensor = tf_tensor;
+  const auto const_t_or = TensorView::New(&const_tf_tensor);
+  ASSERT_TRUE(const_t_or.ok()) << const_t_or.status();
+  const auto& const_t = const_t_or.value();
+
+  EXPECT_THAT(const_t.Data<::tensorflow::tstring>(),
+              ::testing::ElementsAre("a", "bc", "def", "g", "hi", ""));
+
+  const char expectation[] = R"(
+[["a" "bc"]
+ ["def" "g"]
+ ["hi" ""]])";
+
+  EXPECT_THAT(tf_tensor.SummarizeValue(10, true),
+              ::testing::Eq(absl::string_view(expectation).substr(1)));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/tflite_op_shim.cc

@@ -0,0 +1,230 @@
+/* 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/kernels/shim/tflite_op_shim.h"
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/absl_check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+#include "tensorflow/lite/kernels/shim/tflite_tensor_view.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace shim {
+
+namespace internal {
+absl::StatusOr<AttrValue> GetAttr(const flexbuffers::Map* attr_map,
+                                  const std::string& attr_name) {
+  const auto value = (*attr_map)[attr_name.data()];
+  if (value.IsNull())
+    return absl::InvalidArgumentError(
+        absl::StrCat("Non-existent attribute: ", attr_name));
+  AttrValue ret;
+  switch (value.GetType()) {
+    case ::flexbuffers::FBT_BOOL: {
+      ret = value.AsBool();
+      break;
+    }
+    case ::flexbuffers::FBT_INT: {
+      ret = static_cast<int64_t>(value.AsInt64());
+      break;
+    }
+    case ::flexbuffers::FBT_FLOAT: {
+      ret = value.AsFloat();
+      break;
+    }
+    case ::flexbuffers::FBT_STRING: {
+      const auto str_val = value.AsString();
+      ret = absl::string_view(str_val.c_str(), str_val.length());
+      break;
+    }
+    default: {
+      return absl::FailedPreconditionError(
+          absl::StrCat("Unsupported type for attribute: ", attr_name,
+                       " with value: ", value.ToString()));
+    }
+  }
+  return ret;
+}
+}  // namespace internal
+
+TfLiteInitContext::TfLiteInitContext(const TfLiteContext* context,
+                                     const flexbuffers::Map* attr_map)
+    : attr_map_(attr_map) {}
+
+absl::StatusOr<AttrValue> TfLiteInitContext::GetAttr(
+    const std::string& attr_name) const {
+  return internal::GetAttr(attr_map_, attr_name);
+}
+
+TfLiteInvokeContext::TfLiteInvokeContext(TfLiteContext* context,
+                                         TfLiteNode* node)
+    : context_(context), node_(node) {}
+
+ConstTensorViewOr TfLiteInvokeContext::GetInput(const int idx) const {
+  // Scope is used to ensure tensor_view string contents are flushed
+  const auto* tflite_tensor = ::tflite::GetInput(context_, node_, idx);
+  if (tflite_tensor == nullptr)
+    return absl::InternalError(
+        absl::StrCat("input tensor is null during invocation. idx: ", idx));
+  SH_ASSIGN_OR_RETURN(const TfLiteTensorView& tensor_view,
+                      TensorView::New(tflite_tensor));
+  return std::make_unique<const TfLiteTensorView>(tensor_view);
+}
+
+TensorViewOr TfLiteInvokeContext::GetOutput(const int idx,
+                                            const Shape& output_shape) const {
+  if (!output_shape.has_value()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("output_shape value should be populated. idx: ", idx));
+  }
+  auto* tflite_tensor = ::tflite::GetOutput(context_, node_, idx);
+  if (tflite_tensor == nullptr)
+    return absl::InternalError(
+        absl::StrCat("output tensor is null during invocation. idx: ", idx));
+  if (tflite_tensor->data.raw == nullptr ||
+      tflite_tensor->allocation_type == kTfLiteDynamic) {
+    // Clear out string tensor so previous values are not copied.
+    if (tflite_tensor->type == kTfLiteString) {
+      tflite::DynamicBuffer buf;
+      buf.WriteToTensor(tflite_tensor, /*new_shape=*/nullptr);
+    }
+    std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>
+        output_shape_array(ShapeToTfLiteShape(output_shape.value()),
+                           ::TfLiteIntArrayFree);
+    if (context_->ResizeTensor(context_, tflite_tensor,
+                               output_shape_array.release()) != kTfLiteOk) {
+      return absl::InternalError(
+          absl::StrCat("Failed to resize output tensor. idx: ", idx));
+    }
+  } else {
+    ABSL_DCHECK(TfLiteShapeToShape(tflite_tensor->dims) == output_shape);
+  }
+  SH_ASSIGN_OR_RETURN(TfLiteTensorView tensor_view,
+                      TensorView::New(tflite_tensor));
+  return std::make_unique<TfLiteTensorView>(std::move(tensor_view));
+}
+
+int TfLiteInvokeContext::NumInputs() const {
+  return ::tflite::NumInputs(node_);
+}
+
+int TfLiteInvokeContext::NumOutputs() const {
+  return ::tflite::NumOutputs(node_);
+}
+
+TfLiteShapeInferenceContext::TfLiteShapeInferenceContext(
+    TfLiteContext* context, TfLiteNode* node, const flexbuffers::Map* attr_map,
+    std::vector<Shape>* inferred_shapes)
+    : context_(context),
+      node_(node),
+      attr_map_(attr_map),
+      inferred_shapes_(inferred_shapes) {}
+
+ShapeOr TfLiteShapeInferenceContext::GetInputShape(const int idx) const {
+  auto* tflite_tensor = ::tflite::GetInput(context_, node_, idx);
+  if (tflite_tensor == nullptr)
+    return absl::InternalError(absl::StrCat(
+        "input tensor is null during shape inference. idx: ", idx));
+  return TfLiteShapeToShape(tflite_tensor->dims);
+}
+
+// A function object to set output shape information from a Shape
+// object
+absl::Status TfLiteShapeInferenceContext::SetOutputShape(const int idx,
+                                                         const Shape& shape) {
+  if (idx >= inferred_shapes_->size()) {
+    return absl::InternalError(absl::StrCat("output idx out of bounds: ", idx,
+                                            " >= ", inferred_shapes_->size()));
+  }
+  (*inferred_shapes_)[idx] = shape;
+  return absl::OkStatus();
+}
+
+// A function object to read input tensor during shape inference as a TensorView
+ConstTensorViewOr TfLiteShapeInferenceContext::GetInputTensor(
+    const int idx) const {
+  const auto* tflite_tensor = ::tflite::GetInput(context_, node_, idx);
+  if (tflite_tensor == nullptr)
+    return absl::InternalError(absl::StrCat(
+        "input tensor is null during shape inference. idx: ", idx));
+  if (::tflite::IsConstantTensor(tflite_tensor)) {
+    SH_ASSIGN_OR_RETURN(const TfLiteTensorView& tensor_view,
+                        TensorView::New(tflite_tensor));
+    return std::make_unique<const TfLiteTensorView>(tensor_view);
+  } else {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "input tensor is unavailable during shape inference. idx: ", idx));
+  }
+}
+
+absl::StatusOr<AttrValue> TfLiteShapeInferenceContext::GetAttr(
+    const std::string& attr_name) const {
+  return internal::GetAttr(attr_map_, attr_name);
+}
+
+int TfLiteShapeInferenceContext::NumInputs() const {
+  return ::tflite::NumInputs(node_);
+}
+
+int TfLiteShapeInferenceContext::NumOutputs() const {
+  return ::tflite::NumOutputs(node_);
+}
+
+TfLiteStatus StatusToTfLiteStatus(TfLiteContext* context,
+                                  const absl::Status& status) {
+  if (status.ok()) return kTfLiteOk;
+  const auto error_string = std::string(status.message());
+  TF_LITE_KERNEL_LOG(context, "error: %s", error_string.c_str());
+  return kTfLiteError;
+}
+
+TfLiteIntArray* ShapeToTfLiteShape(const std::vector<int>& shape) {
+  TfLiteIntArray* tflite_shape = TfLiteIntArrayCreate(shape.size());
+  // TfLiteIntArray has a data array inside which is not a pointer so there's a
+  // need for copy
+  if (!shape.empty()) {
+    std::memcpy(tflite_shape->data, shape.data(), sizeof(int) * shape.size());
+  }
+  return tflite_shape;
+}
+
+// Converts an int array representing shape in TFLite to Shape.
+Shape TfLiteShapeToShape(const TfLiteIntArray* tflite_shape) {
+  Shape ret(tflite_shape->size);
+  std::memcpy(ret->data(), tflite_shape->data,
+              sizeof(int) * tflite_shape->size);
+  return ret;
+}
+
+}  // namespace shim
+}  // namespace tflite
+
+// Dummy comment to trigger Tricorder

tensorflow/lite/kernels/shim/tflite_op_shim.h

@@ -0,0 +1,225 @@
+/* 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_KERNELS_SHIM_TFLITE_OP_SHIM_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_OP_SHIM_H_
+
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/shape.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+#include "tensorflow/lite/kernels/shim/tflite_tensor_view.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+// This file contains the TFLite adapter. That is, it takes a `OpKernelShim`
+// class and provides a TFLite kernel out of it.
+
+namespace tflite {
+namespace shim {
+
+// TfLite implementation of the methods during an op kernel initialization
+class TfLiteInitContext : public InitContext<TfLiteInitContext> {
+ public:
+  TfLiteInitContext(const TfLiteContext* context,
+                    const flexbuffers::Map* attr_map);
+  // Read a given attribute
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const;
+
+ private:
+  const flexbuffers::Map* attr_map_;
+};
+
+// TfLite implementation of the methods during an op kernel invocation
+class TfLiteInvokeContext : public InvokeContext<TfLiteInvokeContext> {
+ public:
+  TfLiteInvokeContext(TfLiteContext* context_, TfLiteNode* node_);
+  // Read an input tensor
+  ConstTensorViewOr GetInput(const int idx) const;
+  // Get a mutable output tensor. For output string tensors, this should only
+  // be called once.
+  TensorViewOr GetOutput(const int idx, const Shape& shape) const;
+  // Number of input tensors
+  int NumInputs() const;
+  // Number of output tensors
+  int NumOutputs() const;
+
+ private:
+  absl::Status AssertShapesEqual(const TfLiteIntArray* dims,
+                                 const Shape& output_shape) const;
+
+  std::string ShapeMismatchErrorMsg(const TfLiteIntArray* actual_shape,
+                                    const Shape& expected_shape) const;
+
+  TfLiteContext* context_;
+  TfLiteNode* node_;
+};
+
+// TfLite implementation of the methods during shape inference
+class TfLiteShapeInferenceContext
+    : public ShapeInferenceContext<TfLiteShapeInferenceContext> {
+ public:
+  TfLiteShapeInferenceContext(TfLiteContext* context, TfLiteNode* node,
+                              const flexbuffers::Map* attr_map,
+                              std::vector<Shape>* inferred_shapes);
+  // Read an input tensor shape
+  ShapeOr GetInputShape(const int idx) const;
+  // Set an output tensor shape
+  absl::Status SetOutputShape(const int idx, const Shape& shape);
+  // Read an input tensor during shape inference
+  ConstTensorViewOr GetInputTensor(const int idx) const;
+  // Read a given attribute
+  absl::StatusOr<AttrValue> GetAttr(const std::string& attr_name) const;
+  // Number of input tensors
+  int NumInputs() const;
+  // Number of output tensors
+  int NumOutputs() const;
+
+ private:
+  TfLiteContext* context_;
+  TfLiteNode* node_;
+  const flexbuffers::Map* attr_map_;
+  std::vector<Shape>* inferred_shapes_;
+};
+
+// Convert the absl::Status to a TfLiteStatus and report the error message.
+TfLiteStatus StatusToTfLiteStatus(TfLiteContext* context,
+                                  const absl::Status& status);
+
+// Converts a vector of dims into an int array for TFLite use.
+TfLiteIntArray* ShapeToTfLiteShape(const std::vector<int>& shape);
+
+// Converts an int array representing shape in TFLite to Shape.
+Shape TfLiteShapeToShape(const TfLiteIntArray* tflite_shape);
+
+// An op kernel base class which is an adapter between an Op implementation
+// (OpKernelShim subclass) and TFLite runtime
+template <template <Runtime, typename...> typename Impl, typename... Ts>
+class TfLiteOpKernel {
+ public:
+  using ImplType = Impl<Runtime::kTfLite, Ts...>;
+
+  // Builds a TfLiteRegistration object to register this with the TfLite runtime
+  static TfLiteRegistration* GetTfLiteRegistration() {
+    static TfLiteRegistration r =
+        TfLiteRegistration{Init, Free, Prepare, Invoke};
+    return &r;
+  }
+
+  // Adds this op kernel to the passed in op resolver
+  static void Add(MutableOpResolver* resolver) {
+    resolver->AddCustom(ImplType::OpName(), GetTfLiteRegistration());
+  }
+
+  // The operation name
+  static const char* OpName() { return ImplType::OpName(); }
+
+ protected:
+  // The data that is stored in node::user_data.
+  struct UserData {
+    UserData(const char* buffer, size_t length) {
+      impl = new ImplType;
+      attr_map = new flexbuffers::Map(
+          flexbuffers::GetRoot(reinterpret_cast<const uint8_t*>(buffer), length)
+              .AsMap());
+    }
+
+    // An instance of OpKernelShim<TF or TFLite>.
+    // This is so that the Invoke(), Prepare(), etc. can call Invoke(),
+    // ShapeInference(), ... on the kernel defined using this library.
+    ImplType* impl = nullptr;
+    // Attribute map for the op kernel.
+    // The map needs to be accessible because the library provides
+    // GetAttr() during ShapeInference() which is called during Prepare(). So
+    // this needs to be accessible at that point.
+    const flexbuffers::Map* attr_map = nullptr;
+
+    ~UserData() {
+      delete impl;
+      delete attr_map;
+    }
+  };
+
+  static void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+    auto* user_data = new UserData(buffer, length);
+    TfLiteInitContext ctx(context, user_data->attr_map);
+    auto status = user_data->impl->Init(&ctx);
+    StatusToTfLiteStatus(context, status);
+    return user_data;
+  }
+
+  static void Free(TfLiteContext* context, void* buffer) {
+    delete static_cast<UserData*>(buffer);
+  }
+
+  // Resizes the Output Tensor to their shape. There are two cases:
+  //
+  // case 1: output shape is known after ShapeInference() was called during
+  //     Prepare()
+  //   ResizeTensor(output_of_shape_inference)
+  // case 2: output shape is not fully defined even after shape inference
+  //   SetTensorToDynamic(...)
+  static TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+    const size_t num_outputs = ::tflite::NumOutputs(node);
+    std::vector<Shape> inferred_output_shapes(num_outputs);
+    const auto* attr_map = static_cast<UserData*>(node->user_data)->attr_map;
+    TfLiteShapeInferenceContext ctx(context, node, attr_map,
+                                    &inferred_output_shapes);
+    auto status = ImplType::ShapeInference(&ctx);
+    TF_LITE_ENSURE_STATUS(StatusToTfLiteStatus(context, status));
+    // Output shapes.
+    for (int output_idx = 0; output_idx < num_outputs; ++output_idx) {
+      TfLiteTensor* output_tensor =
+          tflite::GetOutput(context, node, output_idx);
+      TF_LITE_ENSURE(context, output_tensor != nullptr);
+      if (inferred_output_shapes[output_idx].FullyDefined()) {
+        // Case: output shape can be inferred during `Prepare`
+        TF_LITE_ENSURE_OK(context,
+                          context->ResizeTensor(
+                              context, output_tensor,
+                              ShapeToTfLiteShape(
+                                  inferred_output_shapes[output_idx].value())));
+      } else {
+        // Case: output shape is dynamic
+        tflite::SetTensorToDynamic(output_tensor);
+      }
+    }
+    return kTfLiteOk;
+  }
+
+  static TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node) {
+    TfLiteInvokeContext ctx(context, node);
+    return StatusToTfLiteStatus(
+        context, static_cast<UserData*>(node->user_data)->impl->Invoke(&ctx));
+  }
+};
+
+template <>
+struct ContextTypeForRuntime<Runtime::kTfLite> {
+  using Init = TfLiteInitContext;
+  using Invoke = TfLiteInvokeContext;
+  using ShapeInference = TfLiteShapeInferenceContext;
+};
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_OP_SHIM_H_

tensorflow/lite/kernels/shim/tflite_op_wrapper.h

@@ -0,0 +1,369 @@
+/* Copyright 2023 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_KERNELS_SHIM_TFLITE_OP_WRAPPER_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_OP_WRAPPER_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/status_macros.h"
+#include "tensorflow/lite/portable_type_to_tflitetype.h"
+
+namespace tflite {
+namespace shim {
+namespace op_wrapper {
+
+using ::tflite::shim::OpKernelShim;
+using ::tflite::shim::Runtime;
+
+// Represents an attribute which can have many types. The first template
+// parameter should be an AttrName, and the packed parameter should be the
+// list of types of the attribute.
+// TODO(b/265879737): When we begin compiling using C++20, the first template
+// parameter should be changed to the template type.
+template <typename N, typename... T>
+struct Attr {
+  const char* Name() const { return N::Name(); }
+};
+
+// Used to store the name of an attribute.
+template <char const* str>
+struct AttrName {
+  static const char* Name() { return str; }
+};
+
+// Object for passing around types.
+template <typename T>
+struct AttrType {
+  using type = T;
+};
+
+// The following constexprs are used to create the variant type which contains
+// the combinations of our templated op. This variant is what is wrapped and
+// ultimately called by the wrapper op.
+//
+// Example:
+//   TmplOp with Attrs {"AType: {bool, float}", "BType: {int32, int64}"};
+// Call:
+//   const char a_type[]("AType"), b_type[]("BType");
+//   VariantOp<Rt,
+//             TmplOp,
+//              Attr<AttrName<a_type>, bool, float>,
+//              Attr<AttrName<b_type>, int32_t, int64_t>> x;
+// Result:
+//   absl::variant<TmplOp<Rt, bool, int32_t>, TmplOp<Rt, bool, int64_t>,
+//                 TmplOp<Rt, float, int32_t>, TmplOp<Rt, float, int64_t>> x;
+
+// Prepends a type onto a tuple.
+template <typename T, typename... Us>
+static constexpr std::tuple<T, Us...> prependTypeInner(T, std::tuple<Us...>);
+
+// Prepend a type on each inner tuple group. This expression unwraps the inner
+// tuples, and the inner expression performs the prepending.
+template <typename T, typename... Us>
+static constexpr auto prependType(T, std::tuple<Us...>)
+    -> std::tuple<decltype(prependTypeInner(std::declval<T>(),
+                                            std::declval<Us>()))...>;
+
+// Base case for recursively processing all combinations of remaining
+// attributes. The result is a tuple containing each type individually.
+template <typename Name, typename... Ts>
+static constexpr std::tuple<std::tuple<Ts>...> getCombinations(
+    Attr<Name, Ts...>);
+
+// Base case for recursively processing all types of an attribute.
+template <typename Name, typename Head, typename... Attrs>
+static constexpr auto getCombinations(Attr<Name, Head>, Attrs...)
+    -> decltype(prependType(std::declval<Head>(),
+                            getCombinations(std::declval<Attrs>()...)));
+
+// Creates a tuple of tuples from a list of Attrute types by recursively
+// popping the first type off the first attribute and prepending it to the
+// combination of other attribute types. This result is then combined with the
+// recursive processing of other types left.
+template <typename Name, typename Head, typename... Tail, typename... Attrs>
+static constexpr auto getCombinations(Attr<Name, Head, Tail...>, Attrs...)
+    -> decltype(std::tuple_cat(
+        prependType(std::declval<Head>(),
+                    getCombinations(std::declval<Attrs>()...)),
+        getCombinations(std::declval<Attr<Name, Tail...>>(),
+                        std::declval<Attrs>()...)));
+
+// Converts a tuple of types into the corresponding op type.
+template <Runtime Rt, template <Runtime, typename...> typename Op,
+          typename... Ts>
+static constexpr Op<Rt, Ts...> convertTuplesToOpsInner(std::tuple<Ts...>);
+
+// Convert a tuple of types into our op with those types. We first need to
+// unwrap the inner tuples, we can then convert each individually in the
+// inner expression and wrap them back up into a tuple.
+template <Runtime Rt, template <Runtime, typename...> typename Op,
+          typename... Ts>
+static constexpr auto convertTuplesToOps(std::tuple<Ts...>) -> std::tuple<
+    decltype(convertTuplesToOpsInner<Rt, Op>(std::declval<Ts>()))...>;
+
+// Convert a tuple of types into a variant of types.
+template <typename... Ts>
+static constexpr std::variant<Ts...> convertTupleToVariant(std::tuple<Ts...>);
+
+// The variant Op type created with TMP. A tuple of tuples containing the
+// attribute combinations is first created. Then each inner tuple is converted
+// into the op types, and finally the outer tuple is converted into a variant.
+// Note, this uses a struct rather than a type alias because of a C++ limitation
+// with template parameter packs not being deduced for aliases.
+template <Runtime Rt, template <Runtime, typename...> typename Op,
+          typename FirstAttr, typename... OtherAttrs>
+struct VariantOp {
+  using type =
+      decltype(convertTupleToVariant(convertTuplesToOps<Rt, Op>(getCombinations(
+          std::declval<FirstAttr>(), std::declval<OtherAttrs>()...))));
+};
+
+// Intermediate object used by the OpWrapper to properly extend OpKernelShim.
+template <Runtime Rt>
+class OpWrapperExtension : public OpKernelShim<OpWrapperExtension, Rt> {};
+
+// Wraps a polymorphic op to be used by TF Lite. At this time, TF Lite does not
+// support TypeConstraints like TensorFlow. This will wrap the op variants
+// and delegate calls to the correctly typed variant when called.
+//
+// Example usage:
+// Given a templated Op `TmplOp` with Attrs:
+//     Attrs {"AType: {bool, float}", "BType: {int32_t, int64_t}"};
+//
+// We can define our type with the following (note that until C++20, these
+// strings cannot be defined inline):
+//
+// const char a_type[]("AType"), b_type[]("BType");
+// template <shim::Runtime Rt>
+// using OpWrapperType = OpWrapper<Rt, TmplOp,
+//     Attr<AttrName<a_type>, bool, float>,
+//     Attr<AttrName<b_type>, int32_t, int64_t>>;
+template <Runtime Rt, template <Runtime, typename...> typename Op,
+          typename... As>
+class OpWrapper : public OpWrapperExtension<Rt> {
+ public:
+  // This variant can be any permutation of the Op and its template params.
+  using TmplOpType = typename VariantOp<Rt, Op, As...>::type;
+  // For static calls, the exact type shouldn't matter, we just need a type.
+  using TmplOpType0 = typename std::variant_alternative<0, TmplOpType>::type;
+
+  using typename OpKernelShim<OpWrapperExtension, Rt>::InitContext;
+  using typename OpKernelShim<OpWrapperExtension, Rt>::InvokeContext;
+  using typename OpKernelShim<OpWrapperExtension, Rt>::ShapeInferenceContext;
+  OpWrapper() = default;
+
+  // For the static methods, they shouldn't change based on the types.
+  static const char* OpName() { return TmplOpType0::OpName(); }
+  static const char* Doc() { return TmplOpType0::Doc(); }
+
+  static std::vector<std::string> Attrs() { return TmplOpType0::Attrs(); }
+  static std::vector<std::string> Inputs() { return TmplOpType0::Inputs(); }
+  static std::vector<std::string> Outputs() { return TmplOpType0::Outputs(); }
+
+  static absl::Status ShapeInference(ShapeInferenceContext* context) {
+    return TmplOpType0::ShapeInference(context);
+  }
+
+  // Creates the correctly typed wrapped object before delegating the Init call
+  // to it. Invoke will also use this variant.
+  absl::Status Init(InitContext* context) {
+    SH_RETURN_IF_ERROR(SetVariantOp<As...>(context));
+
+    return std::visit(
+        [context](auto&& op) -> absl::Status { return op.Init(context); },
+        *op_);
+  }
+
+  // Call Invoke on the created wrapped object.
+  absl::Status Invoke(InvokeContext* context) {
+    return std::visit(
+        [context](auto&& op) -> absl::Status { return op.Invoke(context); },
+        *op_);
+  }
+
+ private:
+  // Sets op_ to the variant type matching the type attributes provided by the
+  // InitContext. Similar to creating the variant type, we recursively
+  // get all combinations of the attributes.
+  template <typename FirstAttr, typename... Attrs>
+  absl::Status SetVariantOp(InitContext* c) {
+    return CombineAttributeTypes(this, c, FirstAttr{}, Attrs{}...);
+  }
+
+  // A simple object to hold Attrutes while we recursively find the
+  // combinations. When called, it will unwrap the stored types to call the
+  // underlying function.
+  // The template parameters are:
+  //   F: Object to wrap which will be another Forwarder object or the OpWrapper
+  //   Name: AttrName of the attribute.
+  //   T: Type of attribute for this combination.
+  template <typename F, typename Name, typename T>
+  struct Forwarder {
+   public:
+    explicit Forwarder(F* f) : inner(f) {}
+
+    template <typename... Args>
+    absl::Status SetOpCombination(Args... args) {
+      return inner->SetOpCombination(Name::Name(), AttrType<T>{}, args...);
+    }
+
+   private:
+    F* inner;
+  };
+
+  // Recursively processes for each combination of attribute types. First,
+  // running over the first attibute and sub-combinations, then running over
+  // the combinations of the remaining types of the first attribute.
+  template <typename F, typename Name, typename Head, typename... Tail,
+            typename... Attrs>
+  absl::Status CombineAttributeTypes(F* obj, InitContext* c,
+                                     Attr<Name, Head, Tail...>, Attrs... rest) {
+    SH_RETURN_IF_ERROR(
+        ApplyAttrType(obj, c, Name{}, AttrType<Head>{}, rest...));
+
+    return CombineAttributeTypes(obj, c, Attr<Name, Tail...>{}, rest...);
+  }
+
+  // Base case for recursively processing types of an attribute.
+  template <typename F, typename Name, typename... Attrs>
+  absl::Status CombineAttributeTypes(F*, InitContext*, Attr<Name>, Attrs...) {
+    return absl::OkStatus();
+  }
+
+  // Saves the names and types of each attribute in the current combination
+  // in a Forwarder object which will ultimately call a typed function.
+  template <typename F, typename Name, typename T, typename Attr,
+            typename... Attrs>
+  absl::Status ApplyAttrType(F* obj, InitContext* c, Name, AttrType<T>, Attr a,
+                             Attrs... rest) {
+    Forwarder<F, Name, T> forwarder(obj);
+
+    return CombineAttributeTypes(&forwarder, c, a, rest...);
+  }
+
+  // Base case for recursively finding combinations of attributes.
+  template <typename F, typename Name, typename T>
+  absl::Status ApplyAttrType(F* obj, InitContext* c, Name, AttrType<T> t) {
+    return obj->SetOpCombination(Name::Name(), t, c);
+  }
+
+  // Checks the attribute types from the context for this particular attribute
+  // type combination. If correct, we set the op variant to this op combo.
+  //
+  // For this, we actually need to overload the functiona nd create a template
+  // for each number of attributes.
+  template <typename T>
+  absl::Status SetOpCombination(std::string Name1, AttrType<T>,
+                                InitContext* context) {
+    int64_t datatype_1;
+    SH_RETURN_IF_ERROR(context->GetAttr(Name1, &datatype_1));
+    if (datatype_1 == typeToTfLiteType<T>()) {
+      this->op_ = std::make_unique<TmplOpType>(Op<Rt, T>());
+    }
+    return absl::OkStatus();
+  }
+
+  template <typename T, typename U>
+  absl::Status SetOpCombination(std::string Name1, AttrType<T>,
+                                std::string Name2, AttrType<U>,
+                                InitContext* context) {
+    int64_t datatype_1, datatype_2;
+    SH_RETURN_IF_ERROR(context->GetAttr(Name1, &datatype_1));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name2, &datatype_2));
+    if (datatype_1 == typeToTfLiteType<T>() &&
+        datatype_2 == typeToTfLiteType<U>()) {
+      this->op_ = std::make_unique<TmplOpType>(Op<Rt, T, U>());
+    }
+    return absl::OkStatus();
+  }
+
+  template <typename T, typename U, typename V>
+  absl::Status SetOpCombination(std::string Name1, AttrType<T>,
+                                std::string Name2, AttrType<U>,
+                                std::string Name3, AttrType<V>,
+                                InitContext* context) {
+    int64_t datatype_1, datatype_2, datatype_3;
+    SH_RETURN_IF_ERROR(context->GetAttr(Name1, &datatype_1));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name2, &datatype_2));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name3, &datatype_3));
+    if (datatype_1 == typeToTfLiteType<T>() &&
+        datatype_2 == typeToTfLiteType<U>() &&
+        datatype_3 == typeToTfLiteType<V>()) {
+      this->op_ = std::make_unique<TmplOpType>(Op<Rt, T, U, V>());
+    }
+    return absl::OkStatus();
+  }
+
+  template <typename T, typename U, typename V, typename W>
+  absl::Status SetOpCombination(std::string Name1, AttrType<T>,
+                                std::string Name2, AttrType<U>,
+                                std::string Name3, AttrType<V>,
+                                std::string Name4, AttrType<W>,
+                                InitContext* context) {
+    int64_t datatype_1, datatype_2, datatype_3, datatype_4;
+    SH_RETURN_IF_ERROR(context->GetAttr(Name1, &datatype_1));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name2, &datatype_2));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name3, &datatype_3));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name4, &datatype_4));
+    if (datatype_1 == typeToTfLiteType<T>() &&
+        datatype_2 == typeToTfLiteType<U>() &&
+        datatype_3 == typeToTfLiteType<V>() &&
+        datatype_4 == typeToTfLiteType<W>()) {
+      this->op_ = std::make_unique<TmplOpType>(Op<Rt, T, U, V, W>());
+    }
+    return absl::OkStatus();
+  }
+
+  template <typename T, typename U, typename V, typename W, typename X>
+  absl::Status SetOpCombination(std::string Name1, AttrType<T>,
+                                std::string Name2, AttrType<U>,
+                                std::string Name3, AttrType<V>,
+                                std::string Name4, AttrType<W>,
+                                std::string Name5, AttrType<X>,
+                                InitContext* context) {
+    int64_t datatype_1, datatype_2, datatype_3, datatype_4, datatype_5;
+    SH_RETURN_IF_ERROR(context->GetAttr(Name1, &datatype_1));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name2, &datatype_2));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name3, &datatype_3));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name4, &datatype_4));
+    SH_RETURN_IF_ERROR(context->GetAttr(Name5, &datatype_5));
+    if (datatype_1 == typeToTfLiteType<T>() &&
+        datatype_2 == typeToTfLiteType<U>() &&
+        datatype_3 == typeToTfLiteType<V>() &&
+        datatype_4 == typeToTfLiteType<W>() &&
+        datatype_5 == typeToTfLiteType<X>()) {
+      this->op_ = std::make_unique<TmplOpType>(Op<Rt, T, U, V, W, X>());
+    }
+    return absl::OkStatus();
+  }
+
+ protected:
+  // The wrapped object variant.
+  std::unique_ptr<TmplOpType> op_;
+};
+
+}  // namespace op_wrapper
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_OP_WRAPPER_H_

tensorflow/lite/kernels/shim/tflite_op_wrapper_test.cc

@@ -0,0 +1,575 @@
+/* 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/kernels/shim/tflite_op_wrapper.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/lite/kernels/shim/op_kernel.h"
+#include "tensorflow/lite/kernels/shim/tflite_op_shim.h"
+
+namespace tflite {
+namespace shim {
+namespace op_wrapper {
+namespace {
+
+#ifndef EXPECT_OK
+#define EXPECT_OK(x) EXPECT_TRUE(x.ok());
+#endif
+
+// Tests the created type of the variant is correct.
+class VariantOpTest : public ::testing::Test {
+ public:
+  // Fake template op to test against.
+  template <shim::Runtime Rt, typename... Ts>
+  class TmplOp {};
+
+  // For checking if variant has a member type
+  template <typename T, typename VARIANT_T>
+  struct isVariantMember;
+
+  template <typename T, typename... ALL_T>
+  struct isVariantMember<T, std::variant<ALL_T...>>
+      : public std::disjunction<std::is_same<T, ALL_T>...> {};
+
+  // Names for parameters
+  static constexpr char kAttrName[] = "AttrName";
+};
+
+TEST_F(VariantOpTest, TestVariantOpCreation_1) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 1);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_2) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, bool>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 2);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_1x1) {
+  using VOp =
+      VariantOp<Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName>, int64_t>,
+                Attr<AttrName<kAttrName>, bool>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 1);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_1x1x1) {
+  using VOp =
+      VariantOp<Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName>, int64_t>,
+                Attr<AttrName<kAttrName>, bool>,
+                Attr<AttrName<kAttrName>, bool>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 1);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, bool>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_2x1) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, float>,
+                        Attr<AttrName<kAttrName>, bool>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 2);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, float, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_1x2) {
+  using VOp =
+      VariantOp<Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName>, int64_t>,
+                Attr<AttrName<kAttrName>, bool, float>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 2);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_2x2) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, int32_t>,
+                        Attr<AttrName<kAttrName>, bool, float>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 4);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_3x3) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, int32_t, int8_t>,
+                        Attr<AttrName<kAttrName>, bool, float, char>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 9);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, char>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, char>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int8_t, bool>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int8_t, float>, VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int8_t, char>, VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_2x2x2) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, int32_t>,
+                        Attr<AttrName<kAttrName>, bool, float>,
+                        Attr<AttrName<kAttrName>, char, int8_t>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 8);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, char>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, int8_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, float, char>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, float, int8_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool, char>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool, int8_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, float, char>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, float, int8_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_2x1x3x1) {
+  using VOp = VariantOp<Runtime::kTfLite, TmplOp,
+                        Attr<AttrName<kAttrName>, int64_t, int32_t>,
+                        Attr<AttrName<kAttrName>, bool>,
+                        Attr<AttrName<kAttrName>, char, int8_t, float>,
+                        Attr<AttrName<kAttrName>, uint16_t>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 6);
+
+  bool b;
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, char, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, int8_t, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int64_t, bool, float, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool, char, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool, int8_t, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+  b = isVariantMember<TmplOp<Runtime::kTfLite, int32_t, bool, float, uint16_t>,
+                      VOp>::value;
+  EXPECT_TRUE(b);
+}
+
+TEST_F(VariantOpTest, TestVariantOpCreation_4x4x6) {
+  using VOp =
+      VariantOp<Runtime::kTfLite, TmplOp,
+                Attr<AttrName<kAttrName>, int64_t, int32_t, int16_t, int8_t>,
+                Attr<AttrName<kAttrName>, int64_t, int32_t, int16_t, int8_t>,
+                Attr<AttrName<kAttrName>, int64_t, int32_t, int16_t, int8_t,
+                     bool, float>>::type;
+
+  EXPECT_EQ(std::variant_size_v<VOp>, 96);
+}
+
+// Tests the correct type of the variant is set given the parameter types of the
+// context in the Init function.
+class SetVariantOpTest : public ::testing::Test {
+ public:
+  // Extend OpWrapper to get access to underlying op variant for testing.
+  template <Runtime Rt, template <Runtime, typename...> typename Op,
+            typename... As>
+  class OpWrapperFriend : public OpWrapper<Rt, Op, As...> {
+   public:
+    using TmplOpType = typename VariantOp<Rt, Op, As...>::type;
+    TmplOpType* GetOp() { return this->op_.get(); }
+  };
+
+  // Fake template op to test against.
+  template <Runtime Rt, typename... Ts>
+  class TmplOp : public OpKernelShim<TmplOp, Rt, Ts...> {
+   public:
+    using typename OpKernelShim<TmplOp, Rt, Ts...>::InitContext;
+    absl::Status Init(InitContext* ctx) { return absl::OkStatus(); }
+  };
+
+  // Fake InitContext used to set the flexbuffer attribute map.
+  class FakeInitContext : public TfLiteInitContext {
+   public:
+    explicit FakeInitContext(const flexbuffers::Map* m)
+        : TfLiteInitContext(nullptr, m) {}
+  };
+
+  // Helper methods for creating a FakeInitContext
+  template <typename T>
+  flexbuffers::Map CreateAttrMap() {
+    fbb_ = std::make_unique<flexbuffers::Builder>();
+    fbb_->Map([&]() {
+      fbb_->Int(kAttrName1, static_cast<int>(typeToTfLiteType<T>()));
+    });
+    fbb_->Finish();
+    return flexbuffers::GetRoot(fbb_->GetBuffer()).AsMap();
+  }
+
+  template <typename T, typename U>
+  flexbuffers::Map CreateAttrMap() {
+    fbb_ = std::make_unique<flexbuffers::Builder>();
+    fbb_->Map([&]() {
+      fbb_->Int(kAttrName1, static_cast<int>(typeToTfLiteType<T>()));
+      fbb_->Int(kAttrName2, static_cast<int>(typeToTfLiteType<U>()));
+    });
+    fbb_->Finish();
+    return flexbuffers::GetRoot(fbb_->GetBuffer()).AsMap();
+  }
+
+  template <typename T, typename U, typename V>
+  flexbuffers::Map CreateAttrMap() {
+    fbb_ = std::make_unique<flexbuffers::Builder>();
+    fbb_->Map([&]() {
+      fbb_->Int(kAttrName1, static_cast<int>(typeToTfLiteType<T>()));
+      fbb_->Int(kAttrName2, static_cast<int>(typeToTfLiteType<U>()));
+      fbb_->Int(kAttrName3, static_cast<int>(typeToTfLiteType<V>()));
+    });
+    fbb_->Finish();
+    return flexbuffers::GetRoot(fbb_->GetBuffer()).AsMap();
+  }
+
+  // Names for parameters
+  static constexpr char kAttrName1[] = "AttrName1";
+  static constexpr char kAttrName2[] = "AttrName2";
+  static constexpr char kAttrName3[] = "AttrName3";
+
+ private:
+  // These must exist for length of test
+  std::unique_ptr<flexbuffers::Builder> fbb_;
+};
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_1) {
+  auto op_wrapper = OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                                    Attr<AttrName<kAttrName1>, bool>>();
+
+  auto map = CreateAttrMap<bool>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_1x1) {
+  auto op_wrapper = OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                                    Attr<AttrName<kAttrName1>, bool>,
+                                    Attr<AttrName<kAttrName2>, int32_t>>();
+
+  auto map = CreateAttrMap<bool, int32_t>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_1x1x1) {
+  auto op_wrapper = OpWrapperFriend<
+      Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName1>, bool>,
+      Attr<AttrName<kAttrName2>, int32_t>, Attr<AttrName<kAttrName3>, float>>();
+
+  auto map = CreateAttrMap<bool, int32_t, float>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b =
+      std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t, float>>(
+          *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_2) {
+  auto op_wrapper =
+      OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                      Attr<AttrName<kAttrName1>, bool, int32_t>>();
+
+  auto map = CreateAttrMap<bool>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b;
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_2x1) {
+  auto op_wrapper = OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                                    Attr<AttrName<kAttrName1>, bool, int32_t>,
+                                    Attr<AttrName<kAttrName2>, float>>();
+
+  auto map = CreateAttrMap<int32_t, float>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b;
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int32_t, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_1x2) {
+  auto op_wrapper =
+      OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                      Attr<AttrName<kAttrName1>, bool>,
+                      Attr<AttrName<kAttrName2>, float, int32_t>>();
+
+  auto map = CreateAttrMap<bool, float>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b;
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_2x2) {
+  auto op_wrapper =
+      OpWrapperFriend<Runtime::kTfLite, TmplOp,
+                      Attr<AttrName<kAttrName1>, bool, int64_t>,
+                      Attr<AttrName<kAttrName2>, float, int32_t>>();
+
+  auto map = CreateAttrMap<bool, float>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b;
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+
+  map = CreateAttrMap<bool, int32_t>();
+  context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+
+  map = CreateAttrMap<int64_t, float>();
+  context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+
+  map = CreateAttrMap<int64_t, int32_t>();
+  context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, bool, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, float>>(
+      *op_wrapper.GetOp());
+  EXPECT_FALSE(b);
+  b = std::holds_alternative<TmplOp<Runtime::kTfLite, int64_t, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_3x3) {
+  auto op_wrapper = OpWrapperFriend<
+      Runtime::kTfLite, TmplOp,
+      Attr<AttrName<kAttrName1>, bool, int64_t, ::tensorflow::tstring>,
+      Attr<AttrName<kAttrName2>, float, int32_t, uint32_t>>();
+
+  auto map = CreateAttrMap<::tensorflow::tstring, int32_t>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b;
+  b = std::holds_alternative<
+      TmplOp<Runtime::kTfLite, ::tensorflow::tstring, int32_t>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_2x2x2) {
+  auto op_wrapper = OpWrapperFriend<
+      Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName1>, bool, int32_t>,
+      Attr<AttrName<kAttrName2>, float, uint32_t>,
+      Attr<AttrName<kAttrName3>, ::tensorflow::tstring, int64_t>>();
+
+  auto map = CreateAttrMap<int32_t, uint32_t, ::tensorflow::tstring>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b = std::holds_alternative<
+      TmplOp<Runtime::kTfLite, int32_t, uint32_t, ::tensorflow::tstring>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_2x1x3) {
+  auto op_wrapper = OpWrapperFriend<
+      Runtime::kTfLite, TmplOp, Attr<AttrName<kAttrName1>, bool, int32_t>,
+      Attr<AttrName<kAttrName2>, float>,
+      Attr<AttrName<kAttrName3>, ::tensorflow::tstring, int64_t, uint32_t>>();
+
+  auto map = CreateAttrMap<int32_t, float, ::tensorflow::tstring>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b = std::holds_alternative<
+      TmplOp<Runtime::kTfLite, int32_t, float, ::tensorflow::tstring>>(
+      *op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+TEST_F(SetVariantOpTest, TestSetVariantOp_4x4x6) {
+  auto op_wrapper = OpWrapperFriend<
+      Runtime::kTfLite, TmplOp,
+      Attr<AttrName<kAttrName1>, bool, int32_t, uint32_t, int8_t>,
+      Attr<AttrName<kAttrName2>, float, int16_t, int32_t, uint32_t>,
+      Attr<AttrName<kAttrName3>, int8_t, uint8_t, int64_t, uint64_t, int32_t,
+           uint32_t>>();
+
+  auto map = CreateAttrMap<int32_t, float, uint32_t>();
+  auto context = FakeInitContext(&map);
+  EXPECT_OK(op_wrapper.Init(&context));
+
+  bool b = std::holds_alternative<
+      TmplOp<Runtime::kTfLite, int32_t, float, uint32_t>>(*op_wrapper.GetOp());
+  EXPECT_TRUE(b);
+}
+
+}  // namespace
+}  // namespace op_wrapper
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/tflite_tensor_view.cc

@@ -0,0 +1,159 @@
+/* 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/kernels/shim/tflite_tensor_view.h"
+
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/variant.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/type_to_tflitetype.h"
+
+// Creates a case statement for the switch() clause given the dtype
+#define CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(TFLITE_DTYPE, CPP_DTYPE) \
+  case TFLITE_DTYPE: {                                          \
+    using DType = typename CPP_DTYPE;                           \
+    return TfLiteTensorView(wrapped_tensor, DType());           \
+  }
+
+#define CASE_FOR_DTYPE(TFLITE_DTYPE) \
+  CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(    \
+      TFLITE_DTYPE, ::tflite::TfLiteTypeToType<TFLITE_DTYPE>::Type)
+
+namespace tflite {
+namespace shim {
+
+TfLiteTensorView::TfLiteTensorView(::TfLiteTensor *wrapped_tensor,
+                                   const ::tensorflow::tstring &dtype)
+    : TensorView(absl::Span<int>(wrapped_tensor->dims->data,
+                                 wrapped_tensor->dims->size),
+                 nullptr, 0, dtype),
+      wrapped_tensor_(wrapped_tensor),
+      const_wrapped_tensor_(wrapped_tensor) {
+  InitForStringDType();
+}
+
+TfLiteTensorView::TfLiteTensorView(const ::TfLiteTensor *wrapped_tensor,
+                                   const ::tensorflow::tstring &dtype)
+    : TensorView(absl::Span<int>(wrapped_tensor->dims->data,
+                                 wrapped_tensor->dims->size),
+                 nullptr, 0, dtype),
+      const_wrapped_tensor_(wrapped_tensor) {
+  InitForStringDType();
+}
+
+TfLiteTensorView::TfLiteTensorView(TfLiteTensorView &&o) noexcept
+    : TensorView(std::move(o)),
+      wrapped_tensor_(o.wrapped_tensor_),
+      const_wrapped_tensor_(o.const_wrapped_tensor_),
+      str_vec_(std::move(o.str_vec_)) {
+}
+
+TfLiteTensorView::TfLiteTensorView(const TfLiteTensorView &o)
+    : TensorView(o),
+      wrapped_tensor_(o.wrapped_tensor_),
+      const_wrapped_tensor_(o.const_wrapped_tensor_),
+      str_vec_(o.str_vec_) {
+}
+
+TfLiteTensorView &TfLiteTensorView::operator=(TfLiteTensorView &&o) noexcept {
+  wrapped_tensor_ = o.wrapped_tensor_;
+  const_wrapped_tensor_ = o.const_wrapped_tensor_;
+  str_vec_ = std::move(o.str_vec_);
+  TensorView::operator=(std::move(o));
+  return *this;
+}
+
+TfLiteTensorView &TfLiteTensorView::operator=(const TfLiteTensorView &o) {
+  if (&o == this) return *this;
+  TensorView::operator=(o);
+  wrapped_tensor_ = o.wrapped_tensor_;
+  const_wrapped_tensor_ = o.const_wrapped_tensor_;
+  str_vec_ = o.str_vec_;
+  return *this;
+}
+
+void TfLiteTensorView::InitForStringDType() {
+  if (str_vec_ == nullptr) {
+    str_vec_ = std::make_shared<StringBuffer>(this);
+  }
+  data_ = absl::Span<::tensorflow::tstring>(str_vec_->buffer);
+}
+
+TfLiteTensorView::StringBuffer::StringBuffer(TfLiteTensorView *t_view)
+    : wrapped_tensor(t_view->wrapped_tensor_) {
+  buffer.resize(NumElements(t_view->shape_));
+  // Read the TfLite string into the buffer
+  const auto const_wrapped_tensor = t_view->const_wrapped_tensor_;
+  std::size_t str_count;
+  if (const_wrapped_tensor->data.raw == nullptr)
+    str_count = 0;
+  else
+    str_count = ::tflite::GetStringCount(const_wrapped_tensor);
+  for (int i = 0; i < str_count; ++i) {
+    const auto str_ref = ::tflite::GetString(const_wrapped_tensor, i);
+    buffer[i].assign_as_view(str_ref.str, str_ref.len);
+  }
+}
+
+TfLiteTensorView::StringBuffer::~StringBuffer() {
+  if (wrapped_tensor == nullptr) return;
+  tflite::DynamicBuffer buf;
+  for (const auto &s : buffer) buf.AddString(s.data(), s.length());
+  buf.WriteToTensor(wrapped_tensor, /*new_shape=*/nullptr);
+}
+
+template <typename TfLiteTensorType>
+absl::StatusOr<
+    typename MatchConstNess<TfLiteTensorType, TfLiteTensorView>::Type>
+TfLiteTensorViewTemplatizedNew(TfLiteTensorType *wrapped_tensor) {
+  switch (wrapped_tensor->type) {
+    CASE_FOR_DTYPE(kTfLiteBool);
+    CASE_FOR_DTYPE(kTfLiteUInt8);
+    CASE_FOR_DTYPE(kTfLiteUInt64);
+    CASE_FOR_DTYPE(kTfLiteInt8);
+    CASE_FOR_DTYPE(kTfLiteInt16);
+    CASE_FOR_DTYPE(kTfLiteInt32);
+    CASE_FOR_DTYPE(kTfLiteInt64);
+    CASE_FOR_DTYPE(kTfLiteFloat32);
+    CASE_FOR_DTYPE(kTfLiteFloat64);
+    // The DType for kTfLiteString is slightly different as we need to use
+    // tensorflow::tstring rather than std::string
+    CASE_FOR_DTYPE_GIVEN_CPP_DTYPE(kTfLiteString, ::tensorflow::tstring);
+    default: {
+      return absl::UnimplementedError(
+          absl::StrCat("Unsupported dtype: ", wrapped_tensor->type));
+    }
+  }
+}
+
+template <>
+absl::StatusOr<TfLiteTensorView> TensorView::New<::TfLiteTensor>(
+    ::TfLiteTensor *wrapped_tensor) {
+  return TfLiteTensorViewTemplatizedNew(wrapped_tensor);
+}
+
+template <>
+absl::StatusOr<const TfLiteTensorView> TensorView::New<const ::TfLiteTensor>(
+    const ::TfLiteTensor *wrapped_tensor) {
+  return TfLiteTensorViewTemplatizedNew(wrapped_tensor);
+}
+
+}  // namespace shim
+}  // namespace tflite

tensorflow/lite/kernels/shim/tflite_tensor_view.h

@@ -0,0 +1,129 @@
+/* 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_KERNELS_SHIM_TFLITE_TENSOR_VIEW_H_
+#define TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_TENSOR_VIEW_H_
+
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/shim/tensor_view.h"
+
+namespace tflite {
+namespace shim {
+
+// A view over a TFLite tensor without taking ownership. It can either be
+// mutable or immutable.
+class TfLiteTensorView : public TensorView {
+ public:
+  // Move constructor
+  TfLiteTensorView(TfLiteTensorView &&o) noexcept;
+  // Copy constructor
+  TfLiteTensorView(const TfLiteTensorView &o);
+  // Move assignment operator
+  TfLiteTensorView &operator=(TfLiteTensorView &&o) noexcept;
+  // Copy assignment operator
+  TfLiteTensorView &operator=(const TfLiteTensorView &);
+
+ protected:
+  // Templated constructor. Since it's not possible to specify the template
+  // argument directly we place a dummy argument of that type so compiler can
+  // deduce the right template parameter
+  template <typename DType>
+  TfLiteTensorView(::TfLiteTensor *wrapped_tensor, const DType &dtype)
+      : TensorView(absl::Span<int>(wrapped_tensor->dims->data,
+                                   wrapped_tensor->dims->size),
+                   wrapped_tensor->data.raw, wrapped_tensor->bytes, dtype),
+        wrapped_tensor_(wrapped_tensor),
+        const_wrapped_tensor_(wrapped_tensor) {}
+
+  // Specialization for string. (this take precedence over the above template)
+  TfLiteTensorView(::TfLiteTensor *wrapped_tensor,
+                   const ::tensorflow::tstring &dtype);
+
+  // Templated constructor for const input.
+  template <typename DType>
+  TfLiteTensorView(const ::TfLiteTensor *wrapped_tensor, const DType &dtype)
+      : TensorView(absl::Span<int>(wrapped_tensor->dims->data,
+                                   wrapped_tensor->dims->size),
+                   wrapped_tensor->data.raw, wrapped_tensor->bytes, dtype),
+        const_wrapped_tensor_(wrapped_tensor) {}
+
+  // Specialization for const string. (this take precedence over the above
+  // template)
+  TfLiteTensorView(const ::TfLiteTensor *wrapped_tensor,
+                   const ::tensorflow::tstring &dtype);
+
+  // Let the factory implementation use private constructors
+  template <typename TfLiteTensorType>
+  friend absl::StatusOr<
+      typename MatchConstNess<TfLiteTensorType, TfLiteTensorView>::Type>
+  TfLiteTensorViewTemplatizedNew(TfLiteTensorType *wrapped_tensor);
+
+  struct StringBuffer {
+    explicit StringBuffer(TfLiteTensorView *t_view);
+    ~StringBuffer();
+
+    // A vector of string as the intermediate shared buffer between
+    // TensorViews
+    std::vector<::tensorflow::tstring> buffer;
+    // The TFLite tensor to which the contents of the buffer is flushed in
+    // dtor
+    ::TfLiteTensor *wrapped_tensor = nullptr;
+  };
+
+  // Initialize the data_ field for string tensors
+  void InitForStringDType();
+
+  // The wrapped TFLiteTensor
+  ::TfLiteTensor *wrapped_tensor_ = nullptr;
+  // A const version of the wrapped TFLiteTensor used when the input is const
+  const ::TfLiteTensor *const_wrapped_tensor_ = nullptr;
+  // A temporary buffer used to expose TfLite strings tensor as Span<tstring>.
+  // This buffer will be flushed and serialized back to the underlying TfLite
+  // string tensor once all the TensorViews over that tensor are destructed.
+  std::shared_ptr<StringBuffer> str_vec_ = nullptr;
+};
+
+// Mapping of ::TfLiteTensor -> TfLiteTensorView
+template <>
+struct TensorViewSubType<::TfLiteTensor> {
+  using Type = TfLiteTensorView;
+};
+
+// Mapping of const ::TfLiteTensor -> const TfLiteTensorView
+template <>
+struct TensorViewSubType<const ::TfLiteTensor> {
+  using Type = const TfLiteTensorView;
+};
+
+// Specialization for TensorView::New()
+template <>
+absl::StatusOr<TfLiteTensorView> TensorView::New<::TfLiteTensor>(
+    ::TfLiteTensor *wrapped_tensor);
+
+// Specialization for TensorView::New()
+template <>
+absl::StatusOr<const TfLiteTensorView> TensorView::New<const ::TfLiteTensor>(
+    const ::TfLiteTensor *wrapped_tensor);
+
+}  // namespace shim
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_SHIM_TFLITE_TENSOR_VIEW_H_

tensorflow/lite/kernels/shim/tflite_tensor_view_test.cc

@@ -0,0 +1,169 @@
+/* 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/kernels/shim/tflite_tensor_view.h"
+
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/shim/test_util.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace shim {
+namespace {
+
+using ::testing::Eq;
+
+TEST(TfLiteTensorW, Bool) {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(1);
+  interpreter.AllocateTensors();
+  auto* tflite_tensor = interpreter.tensor(0);
+  ReallocDynamicTensor<bool>({3, 2}, tflite_tensor);
+  tflite_tensor->name = "test_bool";
+  auto owned_tflite_tensor = UniqueTfLiteTensor(tflite_tensor);
+
+  // Test move assignment
+  auto t_premove_or = TensorView::New(tflite_tensor);
+  ASSERT_TRUE(t_premove_or.ok()) << t_premove_or.status();
+  auto t = std::move(t_premove_or.value());
+
+  auto data = t.Data<bool>();
+  for (int32_t i = 0; i < 3 * 2; ++i) data[i] = (i % 5 == 0);
+
+  ASSERT_THAT(TfliteTensorDebugString(tflite_tensor),
+              Eq("[[1, 0], [0, 0], [0, 1]]"));
+}
+
+template <typename IntType>
+void IntTest() {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(1);
+  interpreter.AllocateTensors();
+  auto* tflite_tensor = interpreter.tensor(0);
+  ReallocDynamicTensor<IntType>({3, 2}, tflite_tensor);
+  tflite_tensor->name = "test_int";
+  auto owned_tflite_tensor = UniqueTfLiteTensor(tflite_tensor);
+
+  // Test move assignment
+  auto t_premove_or = TensorView::New(tflite_tensor);
+  ASSERT_TRUE(t_premove_or.ok()) << t_premove_or.status();
+  auto t = std::move(t_premove_or.value());
+
+  auto data = t.Data<IntType>();
+  for (int32_t i = 0; i < 3 * 2; ++i) data[i] = i;
+
+  ASSERT_THAT(TfliteTensorDebugString(tflite_tensor),
+              Eq("[[0, 1], [2, 3], [4, 5]]"));
+}
+
+TEST(TfLiteTensorW, Int8) { IntTest<int8_t>(); }
+TEST(TfLiteTensorW, UInt8) { IntTest<uint8_t>(); }
+TEST(TfLiteTensorW, Int16) { IntTest<int16_t>(); }
+TEST(TfLiteTensorW, Int32) { IntTest<int32_t>(); }
+TEST(TfLiteTensorW, Int64) { IntTest<int64_t>(); }
+
+template <typename FloatType>
+void FloatTest() {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(1);
+  interpreter.AllocateTensors();
+  auto* tflite_tensor = interpreter.tensor(0);
+  ReallocDynamicTensor<FloatType>({3, 2}, tflite_tensor);
+  tflite_tensor->name = "test_float";
+  auto owned_tflite_tensor = UniqueTfLiteTensor(tflite_tensor);
+
+  auto t_or = TensorView::New(tflite_tensor);
+  ASSERT_TRUE(t_or.ok()) << t_or.status();
+  auto& t = t_or.value();
+
+  auto data = t.Data<FloatType>();
+  for (int32_t i = 0; i < 3 * 2; ++i) data[i] = static_cast<FloatType>(i) / 2.;
+
+  ASSERT_THAT(TfliteTensorDebugString(tflite_tensor),
+              Eq("[[0, 0.5], [1, 1.5], [2, 2.5]]"));
+}
+
+TEST(TfLiteTensorW, Float) { FloatTest<float>(); }
+TEST(TfLiteTensorW, Double) { FloatTest<double>(); }
+
+TEST(TfLiteTensorW, Str) {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(1);
+  interpreter.AllocateTensors();
+  auto* tflite_tensor = interpreter.tensor(0);
+  ReallocDynamicTensor<std::string>({3, 2}, tflite_tensor);
+  tflite_tensor->name = "test_str";
+  auto owned_tflite_tensor = UniqueTfLiteTensor(tflite_tensor);
+
+  {
+    auto t_or = TensorView::New(tflite_tensor);
+    ASSERT_TRUE(t_or.ok()) << t_or.status();
+    auto& t = t_or.value();
+    auto t_mat = t.As<::tensorflow::tstring, 2>();
+    t.Data<::tensorflow::tstring>()[0] = "a";
+    t.Data<::tensorflow::tstring>()[1] = "bc";
+    t_mat(1, 0) = "def";
+    t.Data<::tensorflow::tstring>()[3] = "g";
+    t.Data<::tensorflow::tstring>()[4] = "";
+    t_mat(2, 1) = "hi";
+  }
+
+  {
+    auto t_or = TensorView::New(tflite_tensor);
+    ASSERT_TRUE(t_or.ok()) << t_or.status();
+    auto& t = t_or.value();
+    EXPECT_THAT(t.Data<::tensorflow::tstring>(),
+                ::testing::ElementsAre("a", "bc", "def", "g", "", "hi"));
+  }
+
+  const auto const_tflite_tensor = tflite_tensor;
+  {
+    const auto t_or = TensorView::New(const_tflite_tensor);
+    ASSERT_TRUE(t_or.ok()) << t_or.status();
+    const auto& t = t_or.value();
+    EXPECT_THAT(t.Data<::tensorflow::tstring>(),
+                ::testing::ElementsAre("a", "bc", "def", "g", "", "hi"));
+  }
+
+  EXPECT_THAT(TfliteTensorDebugString(tflite_tensor),
+              Eq("[[a, bc], [def, g], [, hi]]"));
+}
+
+TEST(TfLiteTensorW, EmptyStr) {
+  ::tflite::Interpreter interpreter;
+  interpreter.AddTensors(1);
+  interpreter.AllocateTensors();
+  auto* tflite_tensor = interpreter.tensor(0);
+  ReallocDynamicTensor<std::string>(/*shape=*/{0}, tflite_tensor);
+  tflite_tensor->name = "test_str";
+  auto owned_tflite_tensor = UniqueTfLiteTensor(tflite_tensor);
+
+  // Placing tensor_view instance in a block to ensure its dtor runs
+  {
+    auto t_or = TensorView::New(tflite_tensor);
+    ASSERT_TRUE(t_or.ok()) << t_or.status();
+  }
+
+  EXPECT_THAT(GetStringCount(tflite_tensor), Eq(0));
+}
+
+}  // namespace
+}  // namespace shim
+}  // namespace tflite