chore: import upstream snapshot with attribution

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