chore: import upstream snapshot with attribution

tensorflow/lite/kernels/CMakeLists.txt

@@ -0,0 +1,341 @@
+#
+# 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
+#
+#      https://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.
+
+find_package(googletest REQUIRED)
+find_package(google_benchmark REQUIRED)
+find_package(re2 REQUIRED)
+
+# Generate the mutable schema_generated.h header for tests.
+set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow")
+set(SCHEMA_GENERATED_ROOT ${TF_SOURCE_DIR}/compiler/mlir/lite/schema)
+set(SCHEMA_FILE ${SCHEMA_GENERATED_ROOT}/schema.fbs)
+set(SCHEMA_GENERATED_PATH ${SCHEMA_GENERATED_ROOT}/mutable)
+set(SCHEMA_GENERATED_FILE ${SCHEMA_GENERATED_PATH}/schema_generated.h)
+
+# Use the util function in flatbuffer to generate the schema header.
+include(BuildFlatBuffers)
+
+# For cross-compilation purposes a natively compiled 'flatc' compiler is required
+if(${CMAKE_CROSSCOMPILING})
+  set(FLATC_PATHS
+      ${TFLITE_HOST_TOOLS_DIR}
+      ${TFLITE_HOST_TOOLS_DIR}/bin
+      ${TFLITE_HOST_TOOLS_DIR}/flatbuffers-flatc/bin
+      )
+  find_program(FLATC_BIN flatc PATHS ${FLATC_PATHS})
+  if(${FLATC_BIN} STREQUAL "FLATC_BIN-NOTFOUND")
+    message(FATAL_ERROR "Natively compiled 'flatc' compiler has not been found in the following\
+    locations: ${FLATC_PATHS}")
+  else()
+    message(STATUS "Pre-built 'flatc' compiler for cross-compilation purposes found: ${FLATC_BIN}")
+    set(FLATBUFFERS_FLATC_EXECUTABLE ${FLATC_BIN})
+  endif()
+else()
+  set(FLATBUFFERS_FLATC_EXECUTABLE ${CMAKE_BINARY_DIR}/flatbuffers-flatc/bin/flatc)
+endif()
+
+set(FLATBUFFERS_FLATC_SCHEMA_EXTRA_ARGS
+  -c
+  --gen-object-api
+  --gen-mutable
+)
+build_flatbuffers(
+  "${SCHEMA_FILE}"
+  ""
+  mutable_schema_file
+  flatbuffers-flatc
+  "${SCHEMA_GENERATED_PATH}"
+  ""
+  ""
+)
+
+set(DELEGATE_PROVIDERS_SUPP
+  ${TFLITE_SOURCE_DIR}/nnapi/sl/SupportLibrary.cc
+  ${TFLITE_SOURCE_DIR}/tools/delegates/delegate_provider.cc
+  ${TFLITE_SOURCE_DIR}/tools/evaluation/utils.cc
+)
+
+set(DELEGATE_PROVIDERS
+  ${DELEGATE_PROVIDERS_SUPP}
+  ${TFLITE_SOURCE_DIR}/tools/delegates/default_execution_provider.cc
+  # List of delegates referenced as options in the tensorflow/lite/CMakeLists.txt
+  ${TFLITE_SOURCE_DIR}/tools/delegates/gpu_delegate_provider.cc
+  ${TFLITE_SOURCE_DIR}/tools/delegates/nnapi_delegate_provider.cc
+  ${TFLITE_SOURCE_DIR}/tools/delegates/xnnpack_delegate_provider.cc
+)
+
+if(TFLITE_ENABLE_EXTERNAL_DELEGATE)
+  list(APPEND DELEGATE_PROVIDERS ${TFLITE_SOURCE_DIR}/tools/delegates/external_delegate_provider.cc)
+endif()
+
+set(TEST_FRAMEWORK_SRC
+  ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_list.cc
+  ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_util.cc
+  ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc
+  ${TFLITE_SOURCE_DIR}/../compiler/mlir/lite/schema/schema_conversion_utils.cc
+  ${TFLITE_SOURCE_DIR}/tools/command_line_flags.cc
+  ${DELEGATE_PROVIDERS}
+  ${TFLITE_SOURCE_DIR}/tools/optimize/model_utils.cc
+  ${TF_SOURCE_DIR}/compiler/mlir/lite/tools/optimize/operator_property.cc
+  ${TF_SOURCE_DIR}/compiler/mlir/tools/optimize/quantization_utils.cc
+  ${TF_SOURCE_DIR}/compiler/mlir/lite/kernels/internal/runtime_shape.cc
+  ${TFLITE_SOURCE_DIR}/tools/optimize/quantization_utils.cc
+  ${TFLITE_SOURCE_DIR}/tools/tool_params.cc
+  ${TFLITE_SOURCE_DIR}/tools/versioning/op_version.cc
+  ${XLA_SOURCE_DIR}/xla/tsl/platform/default/env_time.cc
+  ${XLA_SOURCE_DIR}/xla/tsl/platform/default/logging.cc
+  internal/test_util.cc
+  acceleration_test_util.cc
+  acceleration_test_util_internal.cc
+  subgraph_test_util.cc
+  test_delegate_providers.cc
+  test_util.cc
+)
+if(NOT _TFLITE_ENABLE_NNAPI)
+  list(APPEND TEST_FRAMEWORK_SRC
+    ${TFLITE_SOURCE_DIR}/nnapi/nnapi_util.cc
+  )
+endif()
+if(NOT TFLITE_ENABLE_GPU)
+  list(APPEND TEST_FRAMEWORK_SRC
+    ${TFLITE_SOURCE_DIR}/tools/versioning/op_signature.cc
+  )
+endif()
+
+set(TEST_FRAMEWORK_OPTIONS "")
+if(TFLITE_ENABLE_XNNPACK)
+  list(APPEND TEST_FRAMEWORK_SRC
+    ${TFLITE_SOURCE_DIR}/tools/delegates/xnnpack_delegate_provider.cc
+    ${TFLITE_SOURCE_DIR}/core/acceleration/configuration/c/xnnpack_plugin.cc)
+else()
+  list(APPEND TEST_FRAMEWORK_OPTIONS "-DTFLITE_WITHOUT_XNNPACK")
+endif()
+
+# Base library to be later linked with the gtest OR gtest_main library
+add_library(tensorflow-lite-test-base ${TEST_FRAMEWORK_SRC})
+target_link_libraries(tensorflow-lite-test-base
+  gmock
+  re2
+  tensorflow-lite
+  benchmark
+)
+add_dependencies(tensorflow-lite-test-base mutable_schema_file)
+target_include_directories(tensorflow-lite-test-base PUBLIC ${SCHEMA_GENERATED_ROOT})
+target_compile_options(tensorflow-lite-test-base PUBLIC ${TEST_FRAMEWORK_OPTIONS})
+
+add_library(tensorflow-lite-test-gtest-main INTERFACE)
+target_link_libraries(tensorflow-lite-test-gtest-main INTERFACE
+  tensorflow-lite-test-base
+  gtest_main
+)
+
+set(TEST_FRAMEWORK_MAIN_SRC test_main.cc)
+add_library(tensorflow-lite-test-external-main ${TEST_FRAMEWORK_MAIN_SRC})
+target_link_libraries(tensorflow-lite-test-external-main
+  # We need the nnapi related symbols later.
+  # Please check:
+  # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tools/delegates/BUILD#L96
+  -Wl,--whole-archive
+  tensorflow-lite-test-base
+  -Wl,--no-whole-archive
+  gtest
+  absl::log
+)
+
+macro(add_kernel_test TEST_SRC TEST_LIB)
+  string(REPLACE "/" "-" TEST_NAME ${TEST_SRC})
+  string(REPLACE ".cc" "" TEST_NAME ${TEST_NAME})
+
+  add_executable(${TEST_NAME} ${TEST_SRC})
+  target_link_libraries(${TEST_NAME}
+    ${TEST_LIB}
+  )
+  add_test(NAME ${TEST_NAME} COMMAND ${TEST_NAME}
+    WORKING_DIRECTORY ${TENSORFLOW_SOURCE_DIR}
+  )
+  set_tests_properties(${TEST_NAME} PROPERTIES LABELS "plain")
+
+  if(_TFLITE_ENABLE_NNAPI OR TFLITE_ENABLE_EXTERNAL_DELEGATE OR TFLITE_ENABLE_XNNPACK)
+    set(DELEGATE_TEST "${TEST_NAME}_delegate")
+    add_test(
+      NAME ${DELEGATE_TEST}
+      COMMAND cmake -DTEST_EXECUTABLE=$<TARGET_FILE:${TEST_NAME}> -P run-tests.cmake
+    )
+    set_tests_properties(${DELEGATE_TEST} PROPERTIES LABELS "delegate")
+  endif()
+endmacro()
+
+# Tests where main() is provided by the file referenced in TEST_FRAMEWORK_MAIN_SRC
+set(TEST_WITH_EXTERNAL_MAIN_LIST
+  internal/averagepool_quantized_test.cc
+  internal/batch_to_space_nd_test.cc
+  internal/conv_per_channel_quantized_16x8_test.cc
+  internal/depthwiseconv_float_test.cc
+  internal/depthwiseconv_per_channel_quantized_16x8_test.cc
+  internal/depthwiseconv_per_channel_quantized_test.cc
+  internal/depthwiseconv_quantized_test.cc
+  internal/log_quantized_test.cc
+  internal/logsoftmax_quantized_test.cc
+  internal/maxpool_quantized_test.cc
+  internal/non_max_suppression_test.cc
+  internal/per_channel_dequantize_test.cc
+  internal/quantization_util_test.cc
+  internal/resize_bilinear_test.cc
+  internal/resize_nearest_neighbor_test.cc
+  internal/softmax_quantized_test.cc
+  internal/strided_slice_logic_test.cc
+  internal/tensor_test.cc
+  internal/tensor_utils_test.cc
+  internal/transpose_utils_test.cc
+  acceleration_test_util_internal_test.cc
+  activations_test.cc
+  add_n_test.cc
+  add_test.cc
+  arg_min_max_test.cc
+  audio_spectrogram_test.cc
+  basic_rnn_test.cc
+  batch_matmul_test.cc
+  batch_to_space_nd_test.cc
+  bidirectional_sequence_lstm_test.cc
+  bidirectional_sequence_rnn_test.cc
+  broadcast_to_test.cc
+  call_once_test.cc
+  cast_test.cc
+  ceil_test.cc
+  comparisons_test.cc
+  complex_support_test.cc
+  concatenation_test.cc
+  conv3d_test.cc
+  conv_mem_test.cc
+  conv_test.cc
+  cumsum_test.cc
+  densify_test.cc
+  depth_to_space_test.cc
+  depthwise_conv_hybrid_test.cc
+  depthwise_conv_test.cc
+  dequantize_test.cc
+  detection_postprocess_test.cc
+  div_test.cc
+  elementwise_test.cc
+  embedding_lookup_sparse_test.cc
+  embedding_lookup_test.cc
+  exp_test.cc
+  expand_dims_test.cc
+  fake_quant_test.cc
+  fill_test.cc
+  floor_div_test.cc
+  floor_mod_test.cc
+  floor_test.cc
+  fully_connected_test.cc
+  gather_nd_test.cc
+  gather_test.cc
+  hashtable_lookup_test.cc
+  hashtable_ops_test.cc
+  if_test.cc
+  l2norm_test.cc
+  local_response_norm_test.cc
+  log_softmax_test.cc
+  logical_test.cc
+  lsh_projection_test.cc
+  lstm_eval_test.cc
+  lstm_test.cc
+  matrix_diag_test.cc
+  matrix_set_diag_test.cc
+  maximum_minimum_test.cc
+  mfcc_test.cc
+  mirror_pad_test.cc
+  mul_test.cc
+  multinomial_test.cc
+  neg_test.cc
+  non_max_suppression_test.cc
+  numeric_verify_test.cc
+  one_hot_test.cc
+  pack_test.cc
+  pad_test.cc
+  pooling_test.cc
+  pow_test.cc
+  quant_basic_lstm_test.cc
+  quantize_test.cc
+  random_ops_test.cc
+  range_test.cc
+  rank_test.cc
+  reduce_test.cc
+  reshape_test.cc
+  resize_bilinear_test.cc
+  resize_nearest_neighbor_test.cc
+  reverse_sequence_test.cc
+  reverse_test.cc
+  rfft2d_test.cc
+  round_test.cc
+  scatter_nd_test.cc
+  segment_sum_test.cc
+  select_test.cc
+  shape_test.cc
+  skip_gram_test.cc
+  slice_test.cc
+  softmax_test.cc
+  space_to_batch_nd_test.cc
+  space_to_depth_test.cc
+  sparse_to_dense_test.cc
+  split_test.cc
+  split_v_test.cc
+  squared_difference_test.cc
+  squeeze_test.cc
+  strided_slice_test.cc
+  sub_test.cc
+  svdf_test.cc
+  # We don't always have the delegate environment(e.g nnapi)
+  # So, skip "test_delegate_providers_test.cc".
+  #test_delegate_providers_test.cc
+  tile_test.cc
+  topk_v2_test.cc
+  transpose_conv_test.cc
+  transpose_test.cc
+  unidirectional_sequence_lstm_test.cc
+  unidirectional_sequence_rnn_test.cc
+  unique_test.cc
+  unpack_test.cc
+  variable_ops_test.cc
+  where_test.cc
+  while_test.cc
+  zeros_like_test.cc
+)
+# Tests where the main() provided by the GoogleTest framework
+set(TEST_WITH_GTEST_MAIN_LIST
+  cpu_backend_gemm_test.cc
+  cpu_backend_threadpool_test.cc
+  eigen_support_test.cc
+  kernel_util_test.cc
+  optional_tensor_test.cc
+  subgraph_test_util_test.cc
+  test_util_test.cc
+)
+
+foreach(test_src IN LISTS TEST_WITH_EXTERNAL_MAIN_LIST)
+  add_kernel_test(${test_src} tensorflow-lite-test-external-main)
+endforeach()
+
+foreach(test_src IN LISTS TEST_WITH_GTEST_MAIN_LIST)
+  add_kernel_test(${test_src} tensorflow-lite-test-gtest-main)
+endforeach()
+
+# Copy the test utility that facilitates cross-compiled kernel tests run with launch arguments
+if(${CMAKE_CROSSCOMPILING})
+  configure_file(
+    ${TFLITE_SOURCE_DIR}/tools/cmake/test_utils/run-tests.cmake
+    ${CMAKE_CURRENT_BINARY_DIR}/run-tests.cmake
+    COPYONLY
+  )
+endif()

tensorflow/lite/kernels/acceleration_test_util.cc

@@ -0,0 +1,36 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/acceleration_test_util.h"
+
+#include <sstream>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "absl/types/optional.h"
+
+namespace tflite {
+
+std::string GetCurrentTestId() {
+  const ::testing::TestInfo* const test_info =
+      ::testing::UnitTest::GetInstance()->current_test_info();
+
+  std::stringstream test_id_stream;
+
+  test_id_stream << test_info->test_suite_name() << "/" << test_info->name();
+
+  return test_id_stream.str();
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/acceleration_test_util.h

@@ -0,0 +1,28 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_H_
+#define TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_H_
+
+#include <string>
+
+namespace tflite {
+
+// Returns the test id to use to retrieve the acceleration configuration
+// in the acceleration allowlist.
+std::string GetCurrentTestId();
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_H_

tensorflow/lite/kernels/acceleration_test_util_internal.cc

@@ -0,0 +1,68 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/acceleration_test_util_internal.h"
+
+#include <ctype.h>
+
+#include <algorithm>
+#include <functional>
+#include <iterator>
+#include <sstream>
+#include <string>
+
+namespace tflite {
+
+void ReadAccelerationConfig(
+    const char* config,
+    const std::function<void(std::string, std::string, bool)>& consumer) {
+  if (config) {
+    std::istringstream istream{config};
+
+    std::string curr_config_line;
+    while (std::getline(istream, curr_config_line)) {
+      // trim whitespaces
+      curr_config_line.erase(
+          curr_config_line.begin(),
+          std::find_if_not(curr_config_line.begin(), curr_config_line.end(),
+                           [](int ch) { return std::isspace(ch); }));
+      // skipping comments and empty lines.
+      if (curr_config_line.empty() || curr_config_line.at(0) == '#') {
+        continue;
+      }
+
+      // split in test id regexp and rest of the config.
+      auto first_sep_pos =
+          std::find(curr_config_line.begin(), curr_config_line.end(), ',');
+
+      bool is_denylist = false;
+      std::string key = curr_config_line;
+      std::string value{};
+      if (first_sep_pos != curr_config_line.end()) {
+        key = std::string(curr_config_line.begin(), first_sep_pos);
+        value = std::string(first_sep_pos + 1, curr_config_line.end());
+      }
+
+      // Regexps starting with '-'' are denylist ones.
+      if (key[0] == '-') {
+        key = key.substr(1);
+        is_denylist = true;
+      }
+
+      consumer(key, value, is_denylist);
+    }
+  }
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/acceleration_test_util_internal.h

@@ -0,0 +1,107 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_INTERNAL_H_
+#define TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_INTERNAL_H_
+
+#include <algorithm>
+#include <atomic>
+#include <functional>
+#include <iterator>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/types/optional.h"
+#include "re2/re2.h"
+
+namespace tflite {
+
+// Reads the acceleration configuration, handles comments and empty lines and
+// the basic data conversion format (split into key, value, recognition of
+// the line being a white or black list entry) and gives the data to the
+// consumer to be inserted into the target collection.
+void ReadAccelerationConfig(
+    const char* config,
+    const std::function<void(std::string, std::string, bool)>& consumer);
+
+template <typename T>
+class ConfigurationEntry {
+ public:
+  ConfigurationEntry(const std::string& test_id_rex, T test_config,
+                     bool is_denylist)
+      : test_id_rex_(new RE2(test_id_rex)),
+        test_config_(test_config),
+        is_denylist_(is_denylist) {}
+
+  bool Matches(const std::string& test_id) const {
+    return RE2::FullMatch(test_id, *test_id_rex_);
+  }
+  bool IsDenylistEntry() const { return is_denylist_; }
+  const T& TestConfig() const { return test_config_; }
+
+  const std::string& TestIdRex() const { return test_id_rex_->pattern(); }
+
+ private:
+  std::unique_ptr<RE2> test_id_rex_;
+  T test_config_;
+  bool is_denylist_;
+};
+
+// Returns the acceleration test configuration for the given test id and
+// the given acceleration configuration type.
+// The configuration type is responsible of providing the test configuration
+// and the parse function to convert configuration lines into configuration
+// objects.
+template <typename T>
+std::optional<T> GetAccelerationTestParam(std::string test_id) {
+  static std::atomic<std::vector<ConfigurationEntry<T>>*> test_config_ptr;
+
+  if (test_config_ptr.load() == nullptr) {
+    auto config = new std::vector<ConfigurationEntry<T>>();
+
+    auto consumer = [&config](std::string key, std::string value_str,
+                              bool is_denylist) mutable {
+      T value = T::ParseConfigurationLine(value_str);
+      config->emplace_back(key, value, is_denylist);
+    };
+
+    ReadAccelerationConfig(T::AccelerationTestConfig(), consumer);
+
+    // Even if it has been already set, it would be just replaced with the
+    // same value, just freeing the old value to avoid leaks
+    auto* prev_val = test_config_ptr.exchange(config);
+    delete prev_val;
+  }
+
+  const std::vector<ConfigurationEntry<T>>* test_config =
+      test_config_ptr.load();
+
+  const auto test_config_iter =
+      std::find_if(test_config->begin(), test_config->end(),
+                   [&test_id](const ConfigurationEntry<T>& elem) {
+                     return elem.Matches(test_id);
+                   });
+  if (test_config_iter != test_config->end() &&
+      !test_config_iter->IsDenylistEntry()) {
+    return std::optional<T>(test_config_iter->TestConfig());
+  } else {
+    return std::optional<T>();
+  }
+}
+
+}  //  namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_ACCELERATION_TEST_UTIL_INTERNAL_H_

tensorflow/lite/kernels/acceleration_test_util_internal_test.cc

@@ -0,0 +1,211 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/acceleration_test_util_internal.h"
+
+#include <functional>
+#include <optional>
+#include <string>
+#include <unordered_map>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+namespace tflite {
+
+using ::testing::Eq;
+using ::testing::Not;
+using ::testing::Test;
+
+struct SimpleConfig {
+ public:
+  static constexpr char kAccelerationTestConfig[] =
+      R"(
+      #test-id,some-other-data
+      test-1,data-1
+      test-2,
+      test-3,data-3
+      test-4.*,data-4
+      -test-5
+      test-6
+      test-7,data-7
+      )";
+
+  static const char* AccelerationTestConfig() {
+    return kAccelerationTestConfig;
+  }
+
+  static SimpleConfig ParseConfigurationLine(const std::string& conf_line) {
+    return {conf_line};
+  }
+
+  std::string value;
+};
+
+class ReadAccelerationConfigTest : public ::testing::Test {
+ public:
+  std::unordered_map<std::string, SimpleConfig> allowlist_;
+  std::unordered_map<std::string, SimpleConfig> denylist_;
+  std::function<void(std::string, std::string, bool)> consumer_ =
+      [this](std::string key, std::string value, bool is_denylist) {
+        if (is_denylist) {
+          denylist_[key] = {value};
+        } else {
+          allowlist_[key] = {value};
+        }
+      };
+};
+
+TEST_F(ReadAccelerationConfigTest, ReadsAKeyOnlyLine) {
+  ReadAccelerationConfig("key", consumer_);
+
+  EXPECT_THAT(allowlist_.find("key"), Not(Eq(allowlist_.end())));
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, ReadsADenylistKeyOnlyLine) {
+  ReadAccelerationConfig("-key", consumer_);
+
+  EXPECT_THAT(denylist_.find("key"), Not(Eq(allowlist_.end())));
+  EXPECT_TRUE(allowlist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, ReadsAKeyValueLine) {
+  ReadAccelerationConfig("key,value", consumer_);
+
+  EXPECT_THAT(allowlist_["key"].value, Eq("value"));
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, ReadsADenyListKeyValueLine) {
+  ReadAccelerationConfig("-key,value", consumer_);
+
+  EXPECT_THAT(denylist_["key"].value, Eq("value"));
+  EXPECT_TRUE(allowlist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, KeysAreLeftTrimmed) {
+  ReadAccelerationConfig("  key,value", consumer_);
+
+  EXPECT_THAT(allowlist_["key"].value, Eq("value"));
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, BlKeysAreLeftTrimmed) {
+  ReadAccelerationConfig("  -key,value", consumer_);
+
+  EXPECT_THAT(denylist_["key"].value, Eq("value"));
+  EXPECT_TRUE(allowlist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, IgnoresCommentedLines) {
+  ReadAccelerationConfig("#key,value", consumer_);
+
+  EXPECT_TRUE(allowlist_.empty());
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, CommentCanHaveTrailingBlanks) {
+  ReadAccelerationConfig("  #key,value", consumer_);
+
+  EXPECT_TRUE(allowlist_.empty());
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, CommentsAreOnlyForTheFullLine) {
+  ReadAccelerationConfig("key,value #comment", consumer_);
+
+  EXPECT_THAT(allowlist_["key"].value, Eq("value #comment"));
+}
+
+TEST_F(ReadAccelerationConfigTest, IgnoresEmptyLines) {
+  ReadAccelerationConfig("", consumer_);
+
+  EXPECT_TRUE(allowlist_.empty());
+  EXPECT_TRUE(denylist_.empty());
+}
+
+TEST_F(ReadAccelerationConfigTest, ParsesMultipleLines) {
+  ReadAccelerationConfig("key1,value1\nkey2,value2\n-key3,value3", consumer_);
+
+  EXPECT_THAT(allowlist_["key1"].value, Eq("value1"));
+  EXPECT_THAT(allowlist_["key2"].value, Eq("value2"));
+  EXPECT_THAT(denylist_["key3"].value, Eq("value3"));
+}
+
+TEST_F(ReadAccelerationConfigTest, ParsesMultipleLinesWithCommentsAndSpaces) {
+  ReadAccelerationConfig("key1,value1\n#comment\n\nkey2,value2", consumer_);
+
+  EXPECT_THAT(allowlist_["key1"].value, Eq("value1"));
+  EXPECT_THAT(allowlist_["key2"].value, Eq("value2"));
+}
+
+TEST_F(ReadAccelerationConfigTest, ParsesMultipleLinesWithMissingConfigValues) {
+  ReadAccelerationConfig("key1\nkey2,value2\nkey3\nkey4,value4", consumer_);
+
+  EXPECT_THAT(allowlist_["key1"].value, Eq(""));
+  EXPECT_THAT(allowlist_["key2"].value, Eq("value2"));
+  EXPECT_THAT(allowlist_["key3"].value, Eq(""));
+  EXPECT_THAT(allowlist_["key4"].value, Eq("value4"));
+}
+
+TEST(GetAccelerationTestParam, LoadsTestConfig) {
+  const auto config_value_maybe =
+      GetAccelerationTestParam<SimpleConfig>("test-3");
+  ASSERT_TRUE(config_value_maybe.has_value());
+  ASSERT_THAT(config_value_maybe.value().value, Eq("data-3"));
+}
+
+TEST(GetAccelerationTestParam, LoadsTestConfigWithEmptyValue) {
+  const auto config_value_maybe =
+      GetAccelerationTestParam<SimpleConfig>("test-2");
+  ASSERT_TRUE(config_value_maybe.has_value());
+  ASSERT_THAT(config_value_maybe.value().value, Eq(""));
+}
+
+TEST(GetAccelerationTestParam, SupportsWildcards) {
+  const auto config_value_maybe =
+      GetAccelerationTestParam<SimpleConfig>("test-41");
+  ASSERT_TRUE(config_value_maybe.has_value());
+  ASSERT_THAT(config_value_maybe.value().value, Eq("data-4"));
+}
+
+TEST(GetAccelerationTestParam, SupportDenylist) {
+  const auto config_value_maybe =
+      GetAccelerationTestParam<SimpleConfig>("test-5");
+  ASSERT_FALSE(config_value_maybe.has_value());
+}
+
+struct UnmatchedSimpleConfig {
+ public:
+  static constexpr const char* kAccelerationTestConfig = nullptr;
+
+  static const char* AccelerationTestConfig() {
+    return kAccelerationTestConfig;
+  }
+
+  static UnmatchedSimpleConfig ParseConfigurationLine(
+      const std::string& conf_line) {
+    return {conf_line};
+  }
+
+  std::string value;
+};
+
+TEST(GetAccelerationTestParam, ReturnEmptyOptionalForNullConfig) {
+  ASSERT_FALSE(
+      GetAccelerationTestParam<UnmatchedSimpleConfig>("test-3").has_value());
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/add.cc

@@ -0,0 +1,505 @@
+/* Copyright 2017 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/internal/optimized/integer_ops/add.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/add.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
+#include "tensorflow/lite/kernels/internal/reference/process_broadcast_shapes.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+
+namespace ops {
+namespace builtin {
+namespace add {
+// This file has three implementation of Add.
+enum KernelType {
+  kReference,
+  kGenericOptimized,  // Neon-free
+  kNeonOptimized,
+};
+
+constexpr int kInputTensor1 = 0;
+constexpr int kInputTensor2 = 1;
+constexpr int kOutputTensor = 0;
+
+struct OpData {
+  // These fields are used in both the general 8-bit -> 8bit quantized path,
+  // and the special 16-bit -> 16bit quantized path
+  int input1_shift;
+  int input2_shift;
+  int32 output_activation_min;
+  int32 output_activation_max;
+
+  // These fields are used only in the general 8-bit -> 8bit quantized path
+  int32 input1_multiplier;
+  int32 input2_multiplier;
+  int32 output_multiplier;
+  int output_shift;
+  int left_shift;
+  int32 input1_offset;
+  int32 input2_offset;
+  int32 output_offset;
+
+  // This parameter is used to indicate whether
+  // parameter scale is power of two.
+  // It is used in 16-bit -> 16-bit quantization.
+  bool pot_scale_int16;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* data = new OpData;
+  return data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  auto* params = reinterpret_cast<TfLiteAddParams*>(node->builtin_data);
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  TF_LITE_ENSURE(context, params);
+
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
+  output->type = input2->type;
+
+  const bool requires_broadcast = !HaveSameShapes(input1, input2);
+
+  TfLiteIntArray* output_size = nullptr;
+  if (requires_broadcast) {
+    TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
+                                   context, input1, input2, &output_size));
+  } else {
+    output_size = TfLiteIntArrayCopy(input1->dims);
+  }
+
+  // 8bit -> 8bit general quantized path, with general rescalings
+  // as well as, int16 -> int16 with general rescalings
+
+  // There are two implementations of ADD operator in case of
+  // 16bit input/output depending on whether the scale parameter is
+  // the power of 2 or not. Currently only implementation for
+  // general case is used, but we need to use another implementation
+  // for older versions.
+  bool general_scale_int16 = false;
+
+  bool input1_scale_is_pot = false;
+  bool input2_scale_is_pot = false;
+  bool output_scale_is_pot = false;
+
+  int input1_scale_log2_rounded{0};
+  int input2_scale_log2_rounded{0};
+  int output_scale_log2_rounded{0};
+
+  bool input_quantized = input1->quantization.type != kTfLiteNoQuantization;
+
+  if (input1->type == kTfLiteInt16 && input2->type == kTfLiteInt16 &&
+      output->type == kTfLiteInt16 && input_quantized) {
+    // In case of int16, quantization is symmetic and
+    // zero point should be zero.
+    TF_LITE_ENSURE_EQ(context, input1->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, input2->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
+
+    general_scale_int16 = !params->pot_scale_int16;
+
+    if (!general_scale_int16) {
+      // Do preparation in the case of the scale parameter is power of 2.
+
+      input1_scale_is_pot =
+          CheckedLog2(input1->params.scale, &input1_scale_log2_rounded);
+
+      input2_scale_is_pot =
+          CheckedLog2(input2->params.scale, &input2_scale_log2_rounded);
+
+      output_scale_is_pot =
+          CheckedLog2(output->params.scale, &output_scale_log2_rounded);
+
+      general_scale_int16 =
+          !input1_scale_is_pot || !input2_scale_is_pot || !output_scale_is_pot;
+    }
+  }
+
+  data->pot_scale_int16 = !general_scale_int16;
+
+  if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8 ||
+      general_scale_int16) {
+    // 8bit -> 8bit general quantized path, with general rescalings
+    // as well as, 16bit -> 16bit with general rescalings
+    data->input1_offset = -input1->params.zero_point;
+    data->input2_offset = -input2->params.zero_point;
+    data->output_offset = output->params.zero_point;
+
+    // The shift is set to 15 for 16-bit and 20 in case of 8-bit, accordingly.
+    // In case of 16-bit we have 65535 << 15 which is less than 1 << 31,
+    // therefore the addition will still fit in a 32 bit accumulator.
+    data->left_shift = general_scale_int16 ? 15 : 20;
+    const double twice_max_input_scale =
+        2 * std::max(input1->params.scale, input2->params.scale);
+    const double real_input1_multiplier =
+        input1->params.scale / twice_max_input_scale;
+    const double real_input2_multiplier =
+        input2->params.scale / twice_max_input_scale;
+    const double real_output_multiplier =
+        twice_max_input_scale /
+        ((1 << data->left_shift) * output->params.scale);
+
+    QuantizeMultiplierSmallerThanOneExp(
+        real_input1_multiplier, &data->input1_multiplier, &data->input1_shift);
+
+    QuantizeMultiplierSmallerThanOneExp(
+        real_input2_multiplier, &data->input2_multiplier, &data->input2_shift);
+
+    QuantizeMultiplierSmallerThanOneExp(
+        real_output_multiplier, &data->output_multiplier, &data->output_shift);
+
+    TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
+        context, params->activation, output, &data->output_activation_min,
+        &data->output_activation_max));
+  } else if (output->type == kTfLiteInt16 && input_quantized) {
+    // 16bit -> 16bit special quantized path, supporting only a rather
+    // narrow case of quantization parameters: zero_points must all be 0
+    // ("symmetric quantization") and scales must be power-of-two (which
+    // we abbreviate as "POT" below). The intended use case for this path
+    // is in LSTM cells, where, due to the constraints of implementing
+    // some of the math in these LSTM cells in fixed-point arithmetic,
+    // we need to have such symmetric, power-of-two quantization
+    // (Fixed-point formats are inherently symmetric, power-of-two).
+    TF_LITE_ENSURE_EQ(context, input1->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, input2->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
+
+    TF_LITE_ENSURE(context, input1_scale_is_pot);
+    TF_LITE_ENSURE(context, input2_scale_is_pot);
+    TF_LITE_ENSURE(context, output_scale_is_pot);
+
+    data->input1_shift = input1_scale_log2_rounded - output_scale_log2_rounded;
+    data->input2_shift = input2_scale_log2_rounded - output_scale_log2_rounded;
+
+    // Shifting of one input is supported. The graph quantization should ensure
+    // that the other input matches the output.
+    TF_LITE_ENSURE(context, data->input1_shift == 0 || data->input2_shift == 0);
+    TF_LITE_ENSURE(context, data->input1_shift <= 0);
+    TF_LITE_ENSURE(context, data->input2_shift <= 0);
+
+    TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
+        context, params->activation, output, &data->output_activation_min,
+        &data->output_activation_max));
+  }
+
+  return context->ResizeTensor(context, output, output_size);
+}
+
+template <KernelType kernel_type>
+void EvalAdd(TfLiteContext* context, TfLiteNode* node, TfLiteAddParams* params,
+             const OpData* data, const TfLiteTensor* input1,
+             const TfLiteTensor* input2, TfLiteTensor* output) {
+  tflite::ArithmeticParams op_params;
+  const bool need_broadcast = optimized_ops::ProcessBroadcastShapes(
+      GetTensorShape(input1), GetTensorShape(input2), &op_params);
+#define TF_LITE_ADD(type, opname, data_type)                             \
+  data_type output_activation_min, output_activation_max;                \
+  CalculateActivationRange(params->activation, &output_activation_min,   \
+                           &output_activation_max);                      \
+  SetActivationParams(output_activation_min, output_activation_max,      \
+                      &op_params);                                       \
+  type::opname(op_params, GetTensorShape(input1),                        \
+               GetTensorData<data_type>(input1), GetTensorShape(input2), \
+               GetTensorData<data_type>(input2), GetTensorShape(output), \
+               GetTensorData<data_type>(output))
+  if (output->type == kTfLiteInt32) {
+    if (kernel_type == kReference) {
+      if (need_broadcast) {
+        TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, int32_t);
+      } else {
+        TF_LITE_ADD(reference_ops, Add, int32_t);
+      }
+    } else {
+      if (need_broadcast) {
+        TF_LITE_ADD(optimized_ops, BroadcastAdd6DSlow, int32_t);
+      } else {
+        TF_LITE_ADD(optimized_ops, Add, int32_t);
+      }
+    }
+  } else if (output->type == kTfLiteInt64) {
+    if (kernel_type == kReference) {
+      if (need_broadcast) {
+        TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, int64_t);
+      } else {
+        TF_LITE_ADD(reference_ops, Add, int64_t);
+      }
+    } else {
+      if (need_broadcast) {
+        TF_LITE_ADD(optimized_ops, BroadcastAdd6DSlow, int64_t);
+      } else {
+        TF_LITE_ADD(optimized_ops, Add, int64_t);
+      }
+    }
+  } else if (output->type == kTfLiteFloat32) {
+    if (kernel_type == kReference) {
+      if (need_broadcast) {
+        TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, float);
+      } else {
+        TF_LITE_ADD(reference_ops, Add, float);
+      }
+    } else {
+      if (need_broadcast) {
+        TF_LITE_ADD(optimized_ops, BroadcastAddDispatch, float);
+      } else {
+        TF_LITE_ADD(optimized_ops, Add, float);
+      }
+    }
+  } else if (output->type == kTfLiteFloat16) {
+    if (need_broadcast) {
+      TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, half);
+    } else {
+      TF_LITE_ADD(reference_ops, Add, half);
+    }
+  } else if (output->type == kTfLiteInt16) {
+    int16_t output_activation_min, output_activation_max;
+    CalculateActivationRange(params->activation, &output_activation_min,
+                             &output_activation_max);
+    SetActivationParams(output_activation_min, output_activation_max,
+                        &op_params);
+    reference_ops::BroadcastAdd6DSlow<int16_t, true>(
+        op_params, GetTensorShape(input1), GetTensorData<int16_t>(input1),
+        GetTensorShape(input2), GetTensorData<int16_t>(input2),
+        GetTensorShape(output), GetTensorData<int16_t>(output));
+  }
+#undef TF_LITE_ADD
+}
+
+template <KernelType kernel_type>
+TfLiteStatus EvalAddQuantized(TfLiteContext* context, TfLiteNode* node,
+                              TfLiteAddParams* params, const OpData* data,
+                              const TfLiteTensor* input1,
+                              const TfLiteTensor* input2,
+                              TfLiteTensor* output) {
+  if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8 ||
+      !data->pot_scale_int16) {
+    tflite::ArithmeticParams op_params;
+    op_params.left_shift = data->left_shift;
+    op_params.input1_offset = data->input1_offset;
+    op_params.input1_multiplier = data->input1_multiplier;
+    op_params.input1_shift = data->input1_shift;
+    op_params.input2_offset = data->input2_offset;
+    op_params.input2_multiplier = data->input2_multiplier;
+    op_params.input2_shift = data->input2_shift;
+    op_params.output_offset = data->output_offset;
+    op_params.output_multiplier = data->output_multiplier;
+    op_params.output_shift = data->output_shift;
+    SetActivationParams(data->output_activation_min,
+                        data->output_activation_max, &op_params);
+    bool need_broadcast = optimized_ops::ProcessBroadcastShapes(
+        GetTensorShape(input1), GetTensorShape(input2), &op_params);
+#define TF_LITE_ADD(type, opname, dtype)                             \
+  type::opname(op_params, GetTensorShape(input1),                    \
+               GetTensorData<dtype>(input1), GetTensorShape(input2), \
+               GetTensorData<dtype>(input2), GetTensorShape(output), \
+               GetTensorData<dtype>(output));
+    if (output->type == kTfLiteInt8) {
+      if (kernel_type == kReference) {
+        if (need_broadcast) {
+          TF_LITE_ADD(reference_integer_ops, BroadcastAdd6DSlow, int8_t);
+        } else {
+          TF_LITE_ADD(reference_integer_ops, Add, int8_t);
+        }
+      } else {
+        if (need_broadcast) {
+          TF_LITE_ADD(optimized_integer_ops, BroadcastAddDispatch, int8_t);
+        } else {
+          TF_LITE_ADD(optimized_integer_ops, Add, int8_t);
+        }
+      }
+    } else if (output->type == kTfLiteInt16) {
+      if (need_broadcast) {
+        TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, int16_t);
+      } else {
+        if (kernel_type == kReference) {
+          reference_ops::Add(
+              op_params, GetTensorShape(input1), GetTensorData<int16_t>(input1),
+              GetTensorShape(input2), GetTensorData<int16_t>(input2),
+              GetTensorShape(output), GetTensorData<int16_t>(output), false);
+        } else {
+          TF_LITE_ADD(optimized_integer_ops, Add, int16_t);
+        }
+      }
+    } else {
+      if (kernel_type == kReference) {
+        if (need_broadcast) {
+          TF_LITE_ADD(reference_ops, BroadcastAdd6DSlow, uint8_t);
+        } else {
+          TF_LITE_ADD(reference_ops, Add, uint8_t);
+        }
+      } else {
+        if (need_broadcast) {
+          TF_LITE_ADD(optimized_ops, BroadcastAddDispatch, uint8_t);
+        } else {
+          TF_LITE_ADD(optimized_ops, Add, uint8_t);
+        }
+      }
+    }
+#undef TF_LITE_ADD
+  } else if (output->type == kTfLiteInt16) {
+    tflite::ArithmeticParams op_params;
+    op_params.input1_shift = data->input1_shift;
+    op_params.input2_shift = data->input2_shift;
+    SetActivationParams(data->output_activation_min,
+                        data->output_activation_max, &op_params);
+#define TF_LITE_ADD(type, opname)                                      \
+  type::opname(op_params, GetTensorShape(input1),                      \
+               GetTensorData<int16_t>(input1), GetTensorShape(input2), \
+               GetTensorData<int16_t>(input2), GetTensorShape(output), \
+               GetTensorData<int16_t>(output))
+    // The quantized version of Add doesn't support activations, so we
+    // always use BroadcastAdd.
+    if (kernel_type == kReference) {
+      TF_LITE_ADD(reference_ops, Add);
+    } else {
+      TF_LITE_ADD(optimized_ops, Add);
+    }
+#undef TF_LITE_ADD
+  }
+
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params = reinterpret_cast<TfLiteAddParams*>(node->builtin_data);
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  if (output->type == kTfLiteFloat32 || output->type == kTfLiteFloat16 ||
+      output->type == kTfLiteInt32 || output->type == kTfLiteInt64 ||
+      (output->quantization.type == kTfLiteNoQuantization &&
+       output->type == kTfLiteInt16)) {
+    EvalAdd<kernel_type>(context, node, params, data, input1, input2, output);
+  } else if (output->type == kTfLiteUInt8 || output->type == kTfLiteInt8 ||
+             output->type == kTfLiteInt16) {
+    TF_LITE_ENSURE_OK(context,
+                      EvalAddQuantized<kernel_type>(context, node, params, data,
+                                                    input1, input2, output));
+  } else {
+    TF_LITE_UNSUPPORTED_TYPE(context, output->type, "Add");
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace add
+
+TfLiteRegistration* Register_ADD_REF() {
+  static TfLiteRegistration r = {
+      add::Init,
+      add::Free,
+      add::Prepare,
+      add::Eval<add::kReference>,
+      /*profiling_string=*/nullptr,
+      /*builtin_code=*/0,
+      /*custom_name=*/nullptr,
+      /*version=*/0,
+      /*registration_external=*/nullptr,
+      /*async_kernel=*/nullptr,
+      kTfLiteInplaceOpInput0Shared | kTfLiteInplaceOpInput1Shared};
+  return &r;
+}
+
+TfLiteRegistration* Register_ADD_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      add::Init,
+      add::Free,
+      add::Prepare,
+      add::Eval<add::kGenericOptimized>,
+      /*profiling_string=*/nullptr,
+      /*builtin_code=*/0,
+      /*custom_name=*/nullptr,
+      /*version=*/0,
+      /*registration_external=*/nullptr,
+      /*async_kernel=*/nullptr,
+      kTfLiteInplaceOpInput0Shared | kTfLiteInplaceOpInput1Shared};
+  return &r;
+}
+
+TfLiteRegistration* Register_ADD_NEON_OPT() {
+  static TfLiteRegistration r = {
+      add::Init,
+      add::Free,
+      add::Prepare,
+      add::Eval<add::kNeonOptimized>,
+      /*profiling_string=*/nullptr,
+      /*builtin_code=*/0,
+      /*custom_name=*/nullptr,
+      /*version=*/0,
+      /*registration_external=*/nullptr,
+      /*async_kernel=*/nullptr,
+      kTfLiteInplaceOpInput0Shared | kTfLiteInplaceOpInput1Shared};
+  return &r;
+}
+
+TfLiteRegistration* Register_ADD() {
+#ifdef USE_NEON
+  return Register_ADD_NEON_OPT();
+#else
+  return Register_ADD_GENERIC_OPT();
+#endif
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/add_n.cc

@@ -0,0 +1,154 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <stdint.h>
+
+#include <algorithm>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace add_n {
+
+constexpr int kInputTensor1 = 0;
+constexpr int kOutputTensor = 0;
+
+struct OpData {
+  // The index of the temporary tensor where temporary accumulations are kept.
+  int scratch_tensor_index;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  context->AddTensors(context, 1, &op_data->scratch_tensor_index);
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  int num_inputs = NumInputs(node);
+  TF_LITE_ENSURE(context, num_inputs >= 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  output->type = input1->type;
+
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(1);
+  node->temporaries->data[0] = op_data->scratch_tensor_index;
+  TfLiteTensor* scratch_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetTemporarySafe(context, node, /*index=*/0, &scratch_tensor));
+  scratch_tensor->type = input1->type;
+  scratch_tensor->allocation_type = kTfLiteArenaRw;
+
+  CpuBackendContext* cpu_backend_context =
+      CpuBackendContext::GetFromContext(context);
+  // Choose the proper number of thread so that:
+  // (1) Each thread gets at least two tensors (1 if we only have 1 input
+  // tensor).
+  // (2) Total thread_count should be bounded by the maximimum allowed threads.
+  // (3) Tensors are distributed evenly across different threads.
+  const int thread_count =
+      std::min(std::max(1, static_cast<int>(num_inputs) / 2),
+               cpu_backend_context->max_num_threads());
+
+  TfLiteIntArray* scratch_shape = TfLiteIntArrayCreate(1);
+  scratch_shape->data[0] = thread_count * NumElements(input1);
+  TF_LITE_ENSURE_OK(
+      context, context->ResizeTensor(context, scratch_tensor, scratch_shape));
+
+  // Check that all input tensors have the same shape and type.
+  for (int i = kInputTensor1 + 1; i < num_inputs; ++i) {
+    const TfLiteTensor* input;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &input));
+    TF_LITE_ENSURE(context, HaveSameShapes(input1, input));
+    TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input->type);
+  }
+
+  // Use the first input node's dimension to be the dimension of the output
+  // node.
+  TfLiteIntArray* input1_dims = input1->dims;
+  TfLiteIntArray* output_dims = TfLiteIntArrayCopy(input1_dims);
+  return context->ResizeTensor(context, output, output_dims);
+}
+
+template <typename T>
+TfLiteStatus EvalAddN(TfLiteContext* context, TfLiteNode* node) {
+  // TODO(haoliang): Initialize all_inputs only once during init.
+  VectorOfTensors<T> all_inputs(*context, *node->inputs);
+  // Safe to use unchecked since caller checks that tensor is valid
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+  int num_inputs = NumInputs(node);
+  // Safe to use unchecked since caller checks that tensor is valid
+  const TfLiteTensor* input1 = GetInput(context, node, kInputTensor1);
+
+  // Fetch backend context and number of threads.
+  CpuBackendContext* cpu_backend_context =
+      CpuBackendContext::GetFromContext(context);
+  TfLiteTensor* scratch_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, 0, &scratch_tensor));
+  optimized_ops::AddN<T>(GetTensorShape(input1), num_inputs, all_inputs.data(),
+                         GetTensorData<T>(output),
+                         GetTensorData<T>(scratch_tensor), cpu_backend_context);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  if (output->type == kTfLiteFloat32) {
+    TF_LITE_ENSURE_OK(context, EvalAddN<float>(context, node));
+  } else if (output->type == kTfLiteInt32) {
+    TF_LITE_ENSURE_OK(context, EvalAddN<int32_t>(context, node));
+  } else {
+    TF_LITE_KERNEL_LOG(context, "AddN only supports FLOAT32|INT32 now, got %s.",
+                       TfLiteTypeGetName(output->type));
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace add_n
+
+TfLiteRegistration* Register_ADD_N() {
+  static TfLiteRegistration r = {add_n::Init, add_n::Free, add_n::Prepare,
+                                 add_n::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/add_n_test.cc

@@ -0,0 +1,68 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <stdint.h>
+
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/add_n_test_common.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(FloatAddNOpModel, AddMultipleTensors) {
+  FloatAddNOpModel m({{TensorType_FLOAT32, {1, 2, 2, 1}},
+                      {TensorType_FLOAT32, {1, 2, 2, 1}},
+                      {TensorType_FLOAT32, {1, 2, 2, 1}}},
+                     {TensorType_FLOAT32, {}});
+  m.PopulateTensor<float>(m.input(0), {-2.0, 0.2, 0.7, 0.8});
+  m.PopulateTensor<float>(m.input(1), {0.1, 0.2, 0.3, 0.5});
+  m.PopulateTensor<float>(m.input(2), {0.5, 0.1, 0.1, 0.2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(),
+              Pointwise(FloatingPointEq(), {-1.4, 0.5, 1.1, 1.5}));
+}
+
+TEST(FloatAddNOpModel, Add2Tensors) {
+  FloatAddNOpModel m(
+      {{TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}},
+      {TensorType_FLOAT32, {}});
+  m.PopulateTensor<float>(m.input(0), {-2.0, 0.2, 0.7, 0.8});
+  m.PopulateTensor<float>(m.input(1), {0.1, 0.2, 0.3, 0.5});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(),
+              Pointwise(FloatingPointEq(), {-1.9, 0.4, 1.0, 1.3}));
+}
+
+TEST(IntegerAddNOpModel, AddMultipleTensors) {
+  IntegerAddNOpModel m({{TensorType_INT32, {1, 2, 2, 1}},
+                        {TensorType_INT32, {1, 2, 2, 1}},
+                        {TensorType_INT32, {1, 2, 2, 1}}},
+                       {TensorType_INT32, {}});
+  m.PopulateTensor<int32_t>(m.input(0), {-20, 2, 7, 8});
+  m.PopulateTensor<int32_t>(m.input(1), {1, 2, 3, 5});
+  m.PopulateTensor<int32_t>(m.input(2), {10, -5, 1, -2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({-9, -1, 11, 11}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/add_n_test_common.h

@@ -0,0 +1,65 @@
+/* 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_ADD_N_TEST_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_ADD_N_TEST_COMMON_H_
+
+#include <vector>
+
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+
+class BaseAddNOpModel : public SingleOpModel {
+ public:
+  BaseAddNOpModel(const std::vector<TensorData>& inputs,
+                  const TensorData& output) {
+    int num_inputs = inputs.size();
+    std::vector<std::vector<int>> input_shapes;
+
+    for (int i = 0; i < num_inputs; ++i) {
+      inputs_.push_back(AddInput(inputs[i]));
+      input_shapes.push_back(GetShape(inputs_[i]));
+    }
+
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_ADD_N, BuiltinOptions_AddNOptions,
+                 CreateAddNOptions(builder_).Union());
+    BuildInterpreter(input_shapes);
+  }
+
+  int input(int i) { return inputs_[i]; }
+
+ protected:
+  std::vector<int> inputs_;
+  int output_;
+};
+
+class FloatAddNOpModel : public BaseAddNOpModel {
+ public:
+  using BaseAddNOpModel::BaseAddNOpModel;
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+};
+
+class IntegerAddNOpModel : public BaseAddNOpModel {
+ public:
+  using BaseAddNOpModel::BaseAddNOpModel;
+
+  std::vector<int32_t> GetOutput() { return ExtractVector<int32_t>(output_); }
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_ADD_N_TEST_COMMON_H_

tensorflow/lite/kernels/arg_min_max.cc

@@ -0,0 +1,290 @@
+/* Copyright 2018 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/internal/reference/arg_min_max.h"
+
+#include <stdint.h>
+
+#include <functional>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace arg_min_max {
+
+constexpr int kInputTensor = 0;
+constexpr int kAxis = 1;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus ResizeOutput(TfLiteContext* context, const TfLiteTensor* input,
+                          const TfLiteTensor* axis, TfLiteTensor* output) {
+  int axis_value;
+  // Retrieve all 8 bytes when axis type is kTfLiteInt64 to avoid data loss.
+  if (axis->type == kTfLiteInt64) {
+    axis_value = static_cast<int>(*GetTensorData<int64_t>(axis));
+  } else {
+    axis_value = *GetTensorData<int>(axis);
+  }
+  if (axis_value < 0) {
+    axis_value += NumDimensions(input);
+  }
+
+  TF_LITE_ENSURE(context, axis_value >= 0);
+  TF_LITE_ENSURE(context, axis_value < NumDimensions(input));
+
+  // Copy the input dimensions to output except the axis dimension.
+  TfLiteIntArray* output_dims = TfLiteIntArrayCreate(NumDimensions(input) - 1);
+  int j = 0;
+  for (int i = 0; i < NumDimensions(input); ++i) {
+    if (i != axis_value) {
+      output_dims->data[j] = SizeOfDimension(input, i);
+      ++j;
+    }
+  }
+  return context->ResizeTensor(context, output, output_dims);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* axis;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kAxis, &axis));
+  // Make sure the axis is only 1 dimension.
+  TF_LITE_ENSURE_EQ(context, NumElements(axis), 1);
+  // Make sure the axis is only either int32 or int64.
+  TF_LITE_ENSURE(context,
+                 axis->type == kTfLiteInt32 || axis->type == kTfLiteInt64);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  auto* params = reinterpret_cast<TfLiteArgMaxParams*>(node->builtin_data);
+  switch (params->output_type) {
+    case kTfLiteInt32:
+      output->type = kTfLiteInt32;
+      break;
+    case kTfLiteInt64:
+      output->type = kTfLiteInt64;
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Unknown index output data type: %d",
+                         params->output_type);
+      return kTfLiteError;
+  }
+
+  // Check conditions for different types.
+  switch (input->type) {
+    case kTfLiteFloat32:
+    case kTfLiteUInt8:
+    case kTfLiteInt8:
+    case kTfLiteInt32:
+    case kTfLiteBool:
+      break;
+
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Unknown input type: %d, only float32, int types "
+                         "and bool are supported",
+                         input->type);
+      return kTfLiteError;
+  }
+
+  TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
+
+  if (IsConstantOrPersistentTensor(axis)) {
+    TF_LITE_ENSURE_STATUS(ResizeOutput(context, input, axis, output));
+  } else {
+    SetTensorToDynamic(output);
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node, bool is_arg_max) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* axis;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kAxis, &axis));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  if (IsDynamicTensor(output)) {
+    TF_LITE_ENSURE_STATUS(ResizeOutput(context, input, axis, output));
+  }
+
+#define TF_LITE_ARG_MIN_MAX(data_type, axis_type, output_type) \
+  optimized_ops::ArgMinMax(                                    \
+      GetTensorShape(input), GetTensorData<data_type>(input),  \
+      GetTensorData<axis_type>(axis), GetTensorShape(output),  \
+      GetTensorData<output_type>(output), is_arg_max)
+  if (axis->type == kTfLiteInt32) {
+    switch (output->type) {
+      case kTfLiteInt32: {
+        switch (input->type) {
+          case kTfLiteFloat32:
+            TF_LITE_ARG_MIN_MAX(float, int32_t, int32_t);
+            break;
+          case kTfLiteUInt8:
+            TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int32_t);
+            break;
+          case kTfLiteInt8:
+            TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int32_t);
+            break;
+          case kTfLiteInt32:
+            TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int32_t);
+            break;
+          case kTfLiteBool:
+            TF_LITE_ARG_MIN_MAX(bool, int32_t, int32_t);
+            break;
+          default:
+            TF_LITE_KERNEL_LOG(context,
+                               "Only float32, uint8, int8, int32 and bool are "
+                               "supported currently, got %s.",
+                               TfLiteTypeGetName(input->type));
+            return kTfLiteError;
+        }
+      } break;
+      case kTfLiteInt64: {
+        switch (input->type) {
+          case kTfLiteFloat32:
+            TF_LITE_ARG_MIN_MAX(float, int32_t, int64_t);
+            break;
+          case kTfLiteUInt8:
+            TF_LITE_ARG_MIN_MAX(uint8_t, int32_t, int64_t);
+            break;
+          case kTfLiteInt8:
+            TF_LITE_ARG_MIN_MAX(int8_t, int32_t, int64_t);
+            break;
+          case kTfLiteInt32:
+            TF_LITE_ARG_MIN_MAX(int32_t, int32_t, int64_t);
+            break;
+          case kTfLiteBool:
+            TF_LITE_ARG_MIN_MAX(bool, int32_t, int64_t);
+            break;
+          default:
+            TF_LITE_KERNEL_LOG(context,
+                               "Only float32, uint8, int8, int32 and bool are "
+                               "supported currently, got %s.",
+                               TfLiteTypeGetName(input->type));
+            return kTfLiteError;
+        }
+      } break;
+      default:
+        TF_LITE_KERNEL_LOG(
+            context, "Only int32 and int64 are supported currently, got %s.",
+            TfLiteTypeGetName(output->type));
+        return kTfLiteError;
+    }
+  } else {
+    switch (output->type) {
+      case kTfLiteInt32: {
+        switch (input->type) {
+          case kTfLiteFloat32:
+            TF_LITE_ARG_MIN_MAX(float, int64_t, int32_t);
+            break;
+          case kTfLiteUInt8:
+            TF_LITE_ARG_MIN_MAX(uint8_t, int64_t, int32_t);
+            break;
+          case kTfLiteInt8:
+            TF_LITE_ARG_MIN_MAX(int8_t, int64_t, int32_t);
+            break;
+          case kTfLiteInt32:
+            TF_LITE_ARG_MIN_MAX(int32_t, int64_t, int32_t);
+            break;
+          case kTfLiteBool:
+            TF_LITE_ARG_MIN_MAX(bool, int64_t, int32_t);
+            break;
+          default:
+            TF_LITE_KERNEL_LOG(context,
+                               "Only float32, uint8, int8, int32 and bool are "
+                               "supported currently, got %s.",
+                               TfLiteTypeGetName(input->type));
+            return kTfLiteError;
+        }
+      } break;
+      case kTfLiteInt64: {
+        switch (input->type) {
+          case kTfLiteFloat32:
+            TF_LITE_ARG_MIN_MAX(float, int64_t, int64_t);
+            break;
+          case kTfLiteUInt8:
+            TF_LITE_ARG_MIN_MAX(uint8_t, int64_t, int64_t);
+            break;
+          case kTfLiteInt8:
+            TF_LITE_ARG_MIN_MAX(int8_t, int64_t, int64_t);
+            break;
+          case kTfLiteInt32:
+            TF_LITE_ARG_MIN_MAX(int32_t, int64_t, int64_t);
+            break;
+          case kTfLiteBool:
+            TF_LITE_ARG_MIN_MAX(bool, int64_t, int64_t);
+            break;
+          default:
+            TF_LITE_KERNEL_LOG(context,
+                               "Only float32, uint8, int8, int32 and bool are "
+                               "supported currently, got %s.",
+                               TfLiteTypeGetName(input->type));
+            return kTfLiteError;
+        }
+      } break;
+      default:
+        TF_LITE_KERNEL_LOG(
+            context, "Only int32 and int64 are supported currently, got %s.",
+            TfLiteTypeGetName(output->type));
+        return kTfLiteError;
+    }
+  }
+#undef TF_LITE_ARG_MIN_MAX
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus ArgMinEval(TfLiteContext* context, TfLiteNode* node) {
+  return Eval(context, node, false);
+}
+
+TfLiteStatus ArgMaxEval(TfLiteContext* context, TfLiteNode* node) {
+  return Eval(context, node, true);
+}
+
+}  // namespace arg_min_max
+
+TfLiteRegistration* Register_ARG_MAX() {
+  static TfLiteRegistration r = {nullptr, nullptr, arg_min_max::Prepare,
+                                 arg_min_max::ArgMaxEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_ARG_MIN() {
+  static TfLiteRegistration r = {nullptr, nullptr, arg_min_max::Prepare,
+                                 arg_min_max::ArgMinEval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/arg_min_max_test.cc

@@ -0,0 +1,349 @@
+/* Copyright 2018 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 <stdint.h>
+
+#include <initializer_list>
+#include <tuple>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class ArgBaseOpModel : public SingleOpModel {
+ public:
+  ArgBaseOpModel(TensorType input_type, int axis_value, TensorType axis_type,
+                 bool constant_axis, TensorType output_type)
+      : axis_value_(axis_value),
+        axis_type_(axis_type),
+        constant_axis_(constant_axis) {
+    input_ = AddInput(input_type);
+    if (constant_axis) {
+      if (axis_type == TensorType_INT64) {
+        axis_ =
+            AddConstInput(axis_type, {static_cast<int64_t>(axis_value)}, {1});
+      } else {
+        axis_ = AddConstInput(axis_type, {axis_value}, {1});
+      }
+    } else {
+      axis_ = AddInput(axis_type);
+    }
+    output_ = AddOutput(output_type);
+  }
+
+  int input() const { return input_; }
+  int axis() const { return axis_; }
+
+  std::vector<int32_t> GetInt32Output() const {
+    return ExtractVector<int32_t>(output_);
+  }
+  std::vector<int64_t> GetInt64Output() const {
+    return ExtractVector<int64_t>(output_);
+  }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  void PopulateAxisIfNeeded() {
+    if (constant_axis_) return;
+    if (axis_type_ == TensorType_INT32) {
+      PopulateTensor<int32_t>(axis(), {axis_value_});
+    } else {
+      PopulateTensor<int64_t>(axis(), {axis_value_});
+    }
+  }
+
+  const int axis_value_;
+  const TensorType axis_type_;
+  const bool constant_axis_;
+
+  int input_;
+  int axis_;
+  int output_;
+};
+
+class ArgMaxOpModel : public ArgBaseOpModel {
+ public:
+  ArgMaxOpModel(std::initializer_list<int> input_shape, TensorType input_type,
+                int axis_value, TensorType axis_type, bool constant_axis,
+                TensorType output_type)
+      : ArgBaseOpModel(input_type, axis_value, axis_type, constant_axis,
+                       output_type) {
+    ArgBaseOpModel::SetBuiltinOp(
+        BuiltinOperator_ARG_MAX, BuiltinOptions_ArgMaxOptions,
+        CreateArgMaxOptions(ArgBaseOpModel::builder_, output_type).Union());
+    ArgBaseOpModel::BuildInterpreter({input_shape, {1}});
+    PopulateAxisIfNeeded();
+  }
+};
+
+class ArgMinOpModel : public ArgBaseOpModel {
+ public:
+  ArgMinOpModel(std::initializer_list<int> input_shape, TensorType input_type,
+                int axis_value, TensorType axis_type, bool constant_axis,
+                TensorType output_type)
+      : ArgBaseOpModel(input_type, axis_value, axis_type, constant_axis,
+                       output_type) {
+    ArgBaseOpModel::SetBuiltinOp(
+        BuiltinOperator_ARG_MIN, BuiltinOptions_ArgMinOptions,
+        CreateArgMinOptions(ArgBaseOpModel::builder_, output_type).Union());
+    ArgBaseOpModel::BuildInterpreter({input_shape, {1}});
+    PopulateAxisIfNeeded();
+  }
+};
+
+// Declare ArgMinMaxOpTest as a parameterized test, where the parameter is a
+// tuple with:
+// - boolean indicating whether to use a constant axis or not.
+// - axis type (TensorType_INT32 or TensorType_INT64)
+// - output type (TensorType_INT32 or TensorType_INT64)
+class ArgMinMaxOpTest : public ::testing::TestWithParam<
+                            std::tuple<bool, TensorType, TensorType>> {
+ public:
+  bool ConstantAxis() const { return std::get<0>(GetParam()); }
+
+  TensorType AxisType() const { return std::get<1>(GetParam()); }
+
+  TensorType OutputType() const { return std::get<2>(GetParam()); }
+
+  void ValidateOutput(const ArgBaseOpModel& model,
+                      const std::vector<int>& expected_output) {
+    if (OutputType() == TensorType_INT32) {
+      EXPECT_THAT(model.GetInt32Output(), ElementsAreArray(expected_output));
+    } else {
+      EXPECT_THAT(model.GetInt64Output(), ElementsAreArray(expected_output));
+    }
+  }
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    ArgMinMaxOpTest, ArgMinMaxOpTest,
+    ::testing::Combine(::testing::Bool(),
+                       ::testing::Values(TensorType_INT32, TensorType_INT64),
+                       ::testing::Values(TensorType_INT32, TensorType_INT64)));
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgFloat) {
+  ArgMaxOpModel model({1, 1, 1, 4}, TensorType_FLOAT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<float>(model.input(), {0.1, 0.9, 0.7, 0.3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgUInt8) {
+  ArgMaxOpModel model({1, 1, 1, 4}, TensorType_UINT8, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<uint8_t>(model.input(), {1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgInt8) {
+  ArgMaxOpModel model({1, 1, 1, 4}, TensorType_INT8, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int8_t>(model.input(), {-1, -9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {2});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgInt) {
+  ArgMaxOpModel model({1, 1, 1, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgBool) {
+  ArgMaxOpModel model({1, 1, 1, 4}, TensorType_BOOL, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<bool>(model.input(), {true, false, false, false});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgMulDimensions) {
+  ArgMaxOpModel model({1, 1, 2, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {3, 1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgNegativeAxis) {
+  ArgMaxOpModel model({1, 1, 2, 4}, TensorType_INT32, -2, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0, 1, 0, 0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 4}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgOutput64) {
+  ArgMaxOpModel model({1, 1, 2, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {10, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0, 1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgFloatLastAxis) {
+  std::vector<float> input{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0};
+  for (int i = 1; i < 10; ++i) {
+    ArgMaxOpModel model({i}, TensorType_FLOAT32, 0, AxisType(), ConstantAxis(),
+                        OutputType());
+    model.PopulateTensor<float>(
+        model.input(), std::vector<float>(input.begin(), input.begin() + i));
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+    ValidateOutput(model, {i - 1});
+  }
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgFloat) {
+  ArgMinOpModel model({1, 1, 1, 4}, TensorType_FLOAT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<float>(model.input(), {0.1, 0.9, 0.7, 0.3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgInt) {
+  ArgMinOpModel model({1, 1, 1, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgBool) {
+  ArgMinOpModel model({1, 1, 1, 4}, TensorType_BOOL, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<bool>(model.input(), {true, false, true, true});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgMulDimensions) {
+  ArgMinOpModel model({1, 1, 2, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0, 0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgNegativeAxis) {
+  ArgMinOpModel model({1, 1, 2, 4}, TensorType_INT32, -2, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {1, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {0, 0, 0, 1});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 4}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgOutput64) {
+  ArgMinOpModel model({1, 1, 2, 4}, TensorType_INT32, 3, AxisType(),
+                      ConstantAxis(), OutputType());
+  model.PopulateTensor<int>(model.input(), {10, 2, 7, 8, 1, 9, 7, 3});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  ValidateOutput(model, {1, 0});
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 2}));
+}
+
+TEST_P(ArgMinMaxOpTest, GetMinArgFloatLastAxis) {
+  std::vector<float> input{1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1};
+  for (int i = 1; i < 10; ++i) {
+    ArgMinOpModel model({i}, TensorType_FLOAT32, 0, AxisType(), ConstantAxis(),
+                        OutputType());
+    model.PopulateTensor<float>(
+        model.input(), std::vector<float>(input.begin(), input.begin() + i));
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+    ValidateOutput(model, {i - 1});
+  }
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgInt8LastAxis) {
+  // Vector size for int8 is 16 elements, so 35 covers two SIMD widths
+  // Plus extras for testing
+  constexpr int INPUT_SIZE = 35;
+  std::vector<int8_t> input;
+  input.reserve(INPUT_SIZE);
+  for (int i = 0; i < INPUT_SIZE; i++) {
+    input.push_back(INPUT_SIZE - i);
+  }
+  for (int i = 1; i < INPUT_SIZE; ++i) {
+    ArgMinOpModel model({i}, TensorType_INT8, 0, AxisType(), ConstantAxis(),
+                        OutputType());
+    model.PopulateTensor<int8_t>(
+        model.input(), std::vector<int8_t>(input.begin(), input.begin() + i));
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+    ValidateOutput(model, {i - 1});
+  }
+}
+
+TEST_P(ArgMinMaxOpTest, GetMaxArgUInt8LastAxis) {
+  // Vector size for int8 is 16 elements, so 35 covers two SIMD widths
+  // Plus extras for testing
+  constexpr int INPUT_SIZE = 35;
+  std::vector<uint8_t> input;
+  input.reserve(INPUT_SIZE);
+  for (unsigned int i = 0; i < INPUT_SIZE; i++) {
+    input.push_back(INPUT_SIZE - i);
+  }
+  for (int i = 1; i < INPUT_SIZE; ++i) {
+    ArgMinOpModel model({i}, TensorType_UINT8, 0, AxisType(), ConstantAxis(),
+                        OutputType());
+    model.PopulateTensor<uint8_t>(
+        model.input(), std::vector<uint8_t>(input.begin(), input.begin() + i));
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+    ValidateOutput(model, {i - 1});
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/assign_variable.cc

@@ -0,0 +1,76 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <stdint.h>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/resource/resource_variable.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace assign_variable {
+
+constexpr int kInputVariableId = 0;
+constexpr int kInputValue = 1;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 0);
+
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputVariableId,
+                                          &input_resource_id_tensor));
+  TF_LITE_ENSURE(context, (input_resource_id_tensor->type == kTfLiteResource ||
+                           input_resource_id_tensor->type == kTfLiteInt32));
+  TF_LITE_ENSURE_EQ(context, NumElements(input_resource_id_tensor), 1);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  Subgraph* subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+
+  const TfLiteTensor* input_resource_id_tensor;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputVariableId,
+                                          &input_resource_id_tensor));
+  const TfLiteTensor* input_value_tensor;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kInputValue, &input_value_tensor));
+
+  int resource_id = input_resource_id_tensor->data.i32[0];
+  auto& resources = subgraph->resources();
+  resource::CreateResourceVariableIfNotAvailable(&resources, resource_id);
+  auto* variable = resource::GetResourceVariable(&resources, resource_id);
+  TF_LITE_ENSURE(context, variable != nullptr);
+  variable->AssignFrom(input_value_tensor);
+
+  return kTfLiteOk;
+}
+
+}  // namespace assign_variable
+
+TfLiteRegistration* Register_ASSIGN_VARIABLE() {
+  static TfLiteRegistration r = {nullptr, nullptr, assign_variable::Prepare,
+                                 assign_variable::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/atan2.cc

@@ -0,0 +1,107 @@
+// Copyright 2021 Google LLC
+//
+// 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 <cmath>
+
+#include "Eigen/Core"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace atan2 {
+
+TfLiteStatus EnsureSameShape(TfLiteContext* context, const TfLiteTensor* a,
+                             const TfLiteTensor* b) {
+  TF_LITE_ENSURE_EQ(context, tflite::NumDimensions(a),
+                    tflite::NumDimensions(b));
+  return TfLiteStatus::kTfLiteOk;
+}
+
+TfLiteStatus Atan2Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
+
+  const TfLiteTensor* input_y = tflite::GetInput(context, node, 0);
+  const TfLiteTensor* input_x = tflite::GetInput(context, node, 1);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+  // Validate size and type constraints
+  TF_LITE_ENSURE_OK(context, EnsureSameShape(context, input_y, input_x));
+  TF_LITE_ENSURE_TYPES_EQ(context, input_y->type, input_x->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_y->type, output->type);
+  TF_LITE_ENSURE(context, input_y->type == kTfLiteFloat32 ||
+                              input_y->type == kTfLiteFloat64 ||
+                              input_y->type == kTfLiteBFloat16 ||
+                              input_y->type == kTfLiteFloat16);
+
+  TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input_y->dims);
+  return context->ResizeTensor(context, output, output_shape);
+}
+
+template <typename Float>
+TfLiteStatus Atan2(const TfLiteTensor* input_y, const TfLiteTensor* input_x,
+                   TfLiteTensor* output) {
+  const Float* data_y = tflite::GetTensorData<Float>(input_y);
+  const Float* data_x = tflite::GetTensorData<Float>(input_x);
+  Float* data_output = tflite::GetTensorData<Float>(output);
+
+  const int64_t num_elements = NumElements(input_y);
+  for (int64_t i = 0; i < num_elements; ++i) {
+    data_output[i] = static_cast<Float>(std::atan2((data_y[i]), data_x[i]));
+  }
+  return TfLiteStatus::kTfLiteOk;
+}
+
+TfLiteStatus Atan2Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input_y = tflite::GetInput(context, node, 0);
+  const TfLiteTensor* input_x = tflite::GetInput(context, node, 1);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+
+  switch (output->type) {
+    case kTfLiteFloat32:
+      TF_LITE_ENSURE_OK(context, Atan2<float>(input_y, input_x, output));
+      break;
+    case kTfLiteFloat64:
+      TF_LITE_ENSURE_OK(context, Atan2<double>(input_y, input_x, output));
+      break;
+    case kTfLiteFloat16:
+      TF_LITE_ENSURE_OK(context, Atan2<Eigen::half>(input_y, input_x, output));
+      break;
+    case kTfLiteBFloat16:
+      TF_LITE_ENSURE_OK(context,
+                        Atan2<Eigen::bfloat16>(input_y, input_x, output));
+      break;
+    default: {
+      TF_LITE_KERNEL_LOG(context, "Unsupported datatype for atan2 output: %s",
+                         TfLiteTypeGetName(output->type));
+      return TfLiteStatus::kTfLiteError;
+    }
+  }
+
+  return TfLiteStatus::kTfLiteOk;
+}
+
+}  // namespace atan2
+
+TfLiteRegistration* Register_ATAN2() {
+  static TfLiteRegistration r = {nullptr, nullptr, atan2::Atan2Prepare,
+                                 atan2::Atan2Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/atan2_custom.cc

@@ -0,0 +1,102 @@
+// Copyright 2021 Google LLC
+//
+// 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 <cmath>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/custom_ops_register.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace atan2 {
+
+TfLiteStatus EnsureSameShape(TfLiteContext* context, const TfLiteTensor* a,
+                             const TfLiteTensor* b) {
+  TF_LITE_ENSURE_EQ(context, tflite::NumDimensions(a),
+                    tflite::NumDimensions(b));
+
+  return TfLiteStatus::kTfLiteOk;
+}
+
+TfLiteStatus Atan2Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, tflite::NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, tflite::NumOutputs(node), 1);
+
+  const TfLiteTensor* input_y = tflite::GetInput(context, node, 0);
+  const TfLiteTensor* input_x = tflite::GetInput(context, node, 1);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+
+  // Validate size and type constraints
+  TF_LITE_ENSURE_OK(context, EnsureSameShape(context, input_y, input_x));
+  TF_LITE_ENSURE_TYPES_EQ(context, input_y->type, input_x->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_y->type, output->type);
+  TF_LITE_ENSURE(context, input_y->type == kTfLiteFloat32 ||
+                              input_y->type == kTfLiteFloat64);
+
+  TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input_y->dims);
+
+  return context->ResizeTensor(context, output, output_shape);
+}
+
+template <typename Float>
+TfLiteStatus Atan2(const TfLiteTensor* input_y, const TfLiteTensor* input_x,
+                   TfLiteTensor* output) {
+  const Float* data_y = tflite::GetTensorData<Float>(input_y);
+  const Float* data_x = tflite::GetTensorData<Float>(input_x);
+  Float* data_output = tflite::GetTensorData<Float>(output);
+
+  const int64_t num_elements = NumElements(input_y);
+  for (int64_t i = 0; i < num_elements; ++i) {
+    data_output[i] = std::atan2(data_y[i], data_x[i]);
+  }
+
+  return TfLiteStatus::kTfLiteOk;
+}
+
+TfLiteStatus Atan2Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input_y = tflite::GetInput(context, node, 0);
+  const TfLiteTensor* input_x = tflite::GetInput(context, node, 1);
+  TfLiteTensor* output = tflite::GetOutput(context, node, 0);
+
+  switch (output->type) {
+    case kTfLiteFloat32:
+      TF_LITE_ENSURE_OK(context, Atan2<float>(input_y, input_x, output));
+      break;
+    case kTfLiteFloat64:
+      TF_LITE_ENSURE_OK(context, Atan2<double>(input_y, input_x, output));
+      break;
+    default: {
+      TF_LITE_KERNEL_LOG(context, "Unsupported datatype for atan2 output: %s",
+                         TfLiteTypeGetName(output->type));
+      return TfLiteStatus::kTfLiteError;
+    }
+  }
+
+  return TfLiteStatus::kTfLiteOk;
+}
+
+}  // namespace atan2
+
+TfLiteRegistration* Register_ATAN2() {
+  static TfLiteRegistration r = {nullptr, nullptr, atan2::Atan2Prepare,
+                                 atan2::Atan2Eval};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/atan2_custom_test.cc

@@ -0,0 +1,108 @@
+// Copyright 2021 Google LLC
+//
+// 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 <cmath>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/custom_ops_register.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace {
+
+template <typename T>
+tflite::TensorType GetTTEnum();
+
+template <>
+tflite::TensorType GetTTEnum<float>() {
+  return tflite::TensorType_FLOAT32;
+}
+
+template <>
+tflite::TensorType GetTTEnum<double>() {
+  return tflite::TensorType_FLOAT64;
+}
+
+class Atan2Model : public tflite::SingleOpModel {
+ public:
+  Atan2Model(tflite::TensorData y, tflite::TensorData x,
+             tflite::TensorData output) {
+    y_ = AddInput(y);
+    x_ = AddInput(x);
+    output_ = AddOutput(output);
+    SetCustomOp("atan2", {}, ops::custom::Register_ATAN2);
+    BuildInterpreter({GetShape(y_), GetShape(x_)});
+  }
+
+  int y_;
+  int x_;
+  int output_;
+
+  template <typename T>
+  std::vector<T> GetOutput(const std::vector<T>& y, const std::vector<T>& x) {
+    PopulateTensor<T>(y_, y);
+    PopulateTensor<T>(x_, x);
+    Invoke();
+    return ExtractVector<T>(output_);
+  }
+};
+
+template <typename Float>
+class Atan2CustomTest : public ::testing::Test {
+ public:
+  using FloatType = Float;
+};
+
+using TestTypes = ::testing::Types<float, double>;
+
+TYPED_TEST_SUITE(Atan2CustomTest, TestTypes);
+
+TYPED_TEST(Atan2CustomTest, TestScalar) {
+  using Float = typename TestFixture::FloatType;
+  tflite::TensorData y = {GetTTEnum<Float>(), {}};
+  tflite::TensorData x = {GetTTEnum<Float>(), {}};
+  tflite::TensorData output = {GetTTEnum<Float>(), {}};
+  Atan2Model m(y, x, output);
+  auto got = m.GetOutput<Float>({0.0}, {0.0});
+  ASSERT_EQ(got.size(), 1);
+  EXPECT_FLOAT_EQ(got[0], 0.0);
+
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({1.0}, {0.0})[0], M_PI / 2);
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({0.0}, {1.0})[0], 0.0);
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({-1.0}, {0.0})[0], -M_PI / 2);
+}
+
+TYPED_TEST(Atan2CustomTest, TestBatch) {
+  using Float = typename TestFixture::FloatType;
+  tflite::TensorData y = {GetTTEnum<Float>(), {4, 2, 1}};
+  tflite::TensorData x = {GetTTEnum<Float>(), {4, 2, 1}};
+  tflite::TensorData output = {GetTTEnum<Float>(), {4, 2, 1}};
+  Atan2Model m(y, x, output);
+
+  std::vector<Float> y_data = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8};
+  std::vector<Float> x_data = {0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1};
+
+  auto got = m.GetOutput<Float>(y_data, x_data);
+
+  ASSERT_EQ(got.size(), 8);
+  for (int i = 0; i < 8; ++i) {
+    EXPECT_FLOAT_EQ(got[i], std::atan2(y_data[i], x_data[i]));
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/atan2_test.cc

@@ -0,0 +1,120 @@
+// Copyright 2021 Google LLC
+//
+// 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 <cmath>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+namespace {
+
+template <typename T>
+tflite::TensorType GetTTEnum();
+template <>
+tflite::TensorType GetTTEnum<float>() {
+  return tflite::TensorType_FLOAT32;
+}
+
+template <>
+tflite::TensorType GetTTEnum<double>() {
+  return tflite::TensorType_FLOAT64;
+}
+
+template <>
+tflite::TensorType GetTTEnum<half>() {
+  return tflite::TensorType_FLOAT16;
+}
+
+template <>
+tflite::TensorType GetTTEnum<Eigen::bfloat16>() {
+  return tflite::TensorType_BFLOAT16;
+}
+
+class Atan2Model : public tflite::SingleOpModel {
+ public:
+  Atan2Model(tflite::TensorData y, tflite::TensorData x,
+             tflite::TensorData output) {
+    y_ = AddInput(y);
+    x_ = AddInput(x);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_ATAN2, BuiltinOptions_NONE, 0);
+    BuildInterpreter({GetShape(y_), GetShape(x_)});
+  }
+
+  template <typename T>
+  std::vector<T> GetOutput(const std::vector<T>& y, const std::vector<T>& x) {
+    PopulateTensor<T>(y_, y);
+    PopulateTensor<T>(x_, x);
+    Invoke();
+    return ExtractVector<T>(output_);
+  }
+
+ private:
+  int y_;
+  int x_;
+  int output_;
+};
+
+template <typename Float>
+class Atan2Test : public ::testing::Test {
+ public:
+  using FloatType = Float;
+};
+
+using TestTypes = ::testing::Types<float, double, half, Eigen::bfloat16>;
+
+TYPED_TEST_SUITE(Atan2Test, TestTypes);
+
+TYPED_TEST(Atan2Test, TestScalar) {
+  using Float = typename TestFixture::FloatType;
+  tflite::TensorData y = {GetTTEnum<Float>(), {}};
+  tflite::TensorData x = {GetTTEnum<Float>(), {}};
+  tflite::TensorData output = {GetTTEnum<Float>(), {}};
+  Atan2Model m(y, x, output);
+
+  auto got = m.GetOutput<Float>({Float(0.0f)}, {Float(0.0f)});
+  ASSERT_EQ(got.size(), 1);
+  EXPECT_FLOAT_EQ(got[0], 0.0);
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({Float(1.0f)}, {Float(0.0f)})[0],
+                  Float(static_cast<float>(M_PI / 2)));
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({Float(0.0f)}, {Float(1.0f)})[0],
+                  Float(0.0f));
+  ASSERT_FLOAT_EQ(m.GetOutput<Float>({Float(-1.0f)}, {Float(0.0f)})[0],
+                  Float(-static_cast<float>(M_PI / 2)));
+}
+
+TYPED_TEST(Atan2Test, TestBatch) {
+  using Float = typename TestFixture::FloatType;
+  tflite::TensorData y = {GetTTEnum<Float>(), {4, 2, 1}};
+  tflite::TensorData x = {GetTTEnum<Float>(), {4, 2, 1}};
+  tflite::TensorData output = {GetTTEnum<Float>(), {4, 2, 1}};
+  Atan2Model m(y, x, output);
+  std::vector<Float> y_data = {Float(0.1f), Float(0.2f), Float(0.3f),
+                               Float(0.4f), Float(0.5f), Float(0.6f),
+                               Float(0.7f), Float(0.8f)};
+  std::vector<Float> x_data = {Float(0.8f), Float(0.7f), Float(0.6f),
+                               Float(0.5f), Float(0.4f), Float(0.3f),
+                               Float(0.2f), Float(0.1f)};
+  auto got = m.GetOutput<Float>(y_data, x_data);
+  ASSERT_EQ(got.size(), 8);
+  for (int i = 0; i < 8; ++i) {
+    EXPECT_FLOAT_EQ(got[i], Float(std::atan2(y_data[i], x_data[i])));
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/audio_spectrogram.cc

@@ -0,0 +1,174 @@
+/* Copyright 2018 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 <math.h>
+#include <stddef.h>
+#include <stdint.h>
+
+#include <vector>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/spectrogram.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace audio_spectrogram {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+enum KernelType {
+  kReference,
+};
+
+typedef struct {
+  int window_size;
+  int stride;
+  bool magnitude_squared;
+  int output_height;
+  internal::Spectrogram* spectrogram;
+} TfLiteAudioSpectrogramParams;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* data = new TfLiteAudioSpectrogramParams;
+
+  const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
+
+  const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap();
+  data->window_size = m["window_size"].AsInt64();
+  data->stride = m["stride"].AsInt64();
+  data->magnitude_squared = m["magnitude_squared"].AsBool();
+
+  data->spectrogram = new internal::Spectrogram;
+
+  return data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  auto* params = reinterpret_cast<TfLiteAudioSpectrogramParams*>(buffer);
+  delete params->spectrogram;
+  delete params;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteAudioSpectrogramParams*>(node->user_data);
+
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input), 2);
+
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
+
+  TF_LITE_ENSURE(context, params->spectrogram->Initialize(params->window_size,
+                                                          params->stride));
+  const int64_t sample_count = input->dims->data[0];
+  const int64_t length_minus_window = (sample_count - params->window_size);
+  if (length_minus_window < 0) {
+    params->output_height = 0;
+  } else {
+    params->output_height = 1 + (length_minus_window / params->stride);
+  }
+  TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
+  output_size->data[0] = input->dims->data[1];
+  output_size->data[1] = params->output_height;
+  output_size->data[2] = params->spectrogram->output_frequency_channels();
+
+  return context->ResizeTensor(context, output, output_size);
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteAudioSpectrogramParams*>(node->user_data);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  const float* input_data = GetTensorData<float>(input);
+
+  const int64_t sample_count = input->dims->data[0];
+  const int64_t channel_count = input->dims->data[1];
+
+  const int64_t output_width = params->spectrogram->output_frequency_channels();
+
+  float* output_flat = GetTensorData<float>(output);
+
+  std::vector<float> input_for_channel(sample_count);
+  for (int64_t channel = 0; channel < channel_count; ++channel) {
+    float* output_slice =
+        output_flat + (channel * params->output_height * output_width);
+    for (int i = 0; i < sample_count; ++i) {
+      input_for_channel[i] = input_data[i * channel_count + channel];
+    }
+    std::vector<std::vector<float>> spectrogram_output;
+    TF_LITE_ENSURE(context, params->spectrogram->Initialize(params->window_size,
+                                                            params->stride));
+    TF_LITE_ENSURE(context,
+                   params->spectrogram->ComputeSquaredMagnitudeSpectrogram(
+                       input_for_channel, &spectrogram_output));
+    TF_LITE_ENSURE_EQ(context, spectrogram_output.size(),
+                      params->output_height);
+    TF_LITE_ENSURE(context, spectrogram_output.empty() ||
+                                (spectrogram_output[0].size() == output_width));
+    for (int row_index = 0; row_index < params->output_height; ++row_index) {
+      const std::vector<float>& spectrogram_row = spectrogram_output[row_index];
+      TF_LITE_ENSURE_EQ(context, spectrogram_row.size(), output_width);
+      float* output_row = output_slice + (row_index * output_width);
+      if (params->magnitude_squared) {
+        for (int i = 0; i < output_width; ++i) {
+          output_row[i] = spectrogram_row[i];
+        }
+      } else {
+        for (int i = 0; i < output_width; ++i) {
+          output_row[i] = sqrtf(spectrogram_row[i]);
+        }
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace audio_spectrogram
+
+TfLiteRegistration* Register_AUDIO_SPECTROGRAM() {
+  static TfLiteRegistration r = {
+      audio_spectrogram::Init, audio_spectrogram::Free,
+      audio_spectrogram::Prepare,
+      audio_spectrogram::Eval<audio_spectrogram::kReference>};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/audio_spectrogram_test.cc

@@ -0,0 +1,113 @@
+/* Copyright 2018 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 <vector>
+
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+TfLiteRegistration* Register_AUDIO_SPECTROGRAM();
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+class BaseAudioSpectrogramOpModel : public SingleOpModel {
+ public:
+  BaseAudioSpectrogramOpModel(const TensorData& input1,
+                              const TensorData& output, int window_size,
+                              int stride, bool magnitude_squared) {
+    input1_ = AddInput(input1);
+    output_ = AddOutput(output);
+
+    flexbuffers::Builder fbb;
+    fbb.Map([&]() {
+      fbb.Int("window_size", window_size);
+      fbb.Int("stride", stride);
+      fbb.Bool("magnitude_squared", magnitude_squared);
+    });
+    fbb.Finish();
+    SetCustomOp("AudioSpectrogram", fbb.GetBuffer(),
+                Register_AUDIO_SPECTROGRAM);
+    BuildInterpreter({GetShape(input1_)});
+  }
+
+  int input1() { return input1_; }
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  int input1_;
+  int output_;
+};
+
+TEST(BaseAudioSpectrogramOpModel, NonSquaredTest) {
+  BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {8, 1}},
+                                {TensorType_FLOAT32, {}}, 8, 1, false);
+  m.PopulateTensor<float>(m.input1(),
+                          {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, 1.0f, 0.0f});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  std::vector<int> output_shape = m.GetOutputShape();
+  EXPECT_EQ(3, output_shape.size());
+  EXPECT_THAT(output_shape, ElementsAre(1, 1, 5));
+
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                 {0.0f, 1.0f, 2.0f, 1.0f, 0.0f}, 1e-3)));
+}
+
+TEST(SpectrogramOpTest, SquaredTest) {
+  BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {8, 1}},
+                                {TensorType_FLOAT32, {}}, 8, 1, true);
+  m.PopulateTensor<float>(m.input1(),
+                          {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f, 1.0f, 0.0f});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  std::vector<int> output_shape = m.GetOutputShape();
+  EXPECT_EQ(3, output_shape.size());
+  EXPECT_THAT(output_shape, ElementsAre(1, 1, 5));
+
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                 {0.f, 1.f, 4.f, 1.f, 0.f}, 1e-3)));
+}
+
+TEST(SpectrogramOpTest, StrideTest) {
+  BaseAudioSpectrogramOpModel m({TensorType_FLOAT32, {10, 1}},
+                                {TensorType_FLOAT32, {}}, 8, 2, true);
+  m.PopulateTensor<float>(m.input1(), {-1.0f, 0.0f, 1.0f, 0.0f, -1.0f, 0.0f,
+                                       1.0f, 0.0f, 1.0f, 0.0f});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  std::vector<int> output_shape = m.GetOutputShape();
+  EXPECT_THAT(output_shape, ElementsAre(1, 2, 5));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                 {0, 1, 4, 1, 0, 1, 2, 1, 2, 1}, 1e-3)));
+}
+
+}  // namespace
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/basic_rnn.cc

@@ -0,0 +1,348 @@
+/* Copyright 2017 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 <cstddef>
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace rnn {
+
+namespace {
+
+struct OpData {
+  int scratch_tensor_index;
+  bool compute_row_sums = false;
+};
+
+}  // namespace
+
+constexpr int kInputTensor = 0;
+constexpr int kWeightsTensor = 1;
+constexpr int kRecurrentWeightsTensor = 2;
+constexpr int kBiasTensor = 3;
+constexpr int kHiddenStateTensor = 4;
+
+// Output tensor.
+constexpr int kOutputTensor = 0;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  context->AddTensors(context, /*tensors_to_add=*/6,
+                      &op_data->scratch_tensor_index);
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  // Check we have all the inputs and outputs we need.
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 5);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* input_weights;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kWeightsTensor, &input_weights));
+  const TfLiteTensor* recurrent_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node, kRecurrentWeightsTensor, &recurrent_weights));
+  const TfLiteTensor* bias;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBiasTensor, &bias));
+  const TfLiteTensor* hidden_state;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kHiddenStateTensor, &hidden_state));
+
+  // Check all the parameters of tensor match within themselves and match the
+  // input configuration.
+  const int batch_size = input->dims->data[0];
+  const int num_units = input_weights->dims->data[0];
+  TF_LITE_ENSURE_EQ(context, input->dims->data[1],
+                    input_weights->dims->data[1]);
+  TF_LITE_ENSURE_EQ(context, input_weights->dims->data[0], bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, recurrent_weights->dims->data[0],
+                    bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, recurrent_weights->dims->data[1],
+                    bias->dims->data[0]);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, input_weights->type,
+                          recurrent_weights->type);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(hidden_state), 2);
+  TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[0], batch_size);
+  TF_LITE_ENSURE_EQ(context, hidden_state->dims->data[1], num_units);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // Resize output.
+  TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(2);
+  output_size_array->data[0] = batch_size;
+  output_size_array->data[1] = num_units;
+  TF_LITE_ENSURE_OK(context,
+                    context->ResizeTensor(context, output, output_size_array));
+
+  const bool is_hybrid = IsHybridOp(input, input_weights);
+
+  // Allocate temporary tensors to store quantized values of input and
+  // hidden_state tensors.
+  if (is_hybrid) {
+    auto* op_data = reinterpret_cast<OpData*>(node->user_data);
+    op_data->compute_row_sums = true;
+    TfLiteIntArrayFree(node->temporaries);
+    node->temporaries = TfLiteIntArrayCreate(6);
+    node->temporaries->data[0] = op_data->scratch_tensor_index;
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/0,
+                                                &input_quantized));
+    input_quantized->type = input_weights->type;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
+      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                       input_quantized_size));
+    }
+    node->temporaries->data[1] = op_data->scratch_tensor_index + 1;
+    TfLiteTensor* hidden_state_quantized;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/1,
+                                                &hidden_state_quantized));
+    hidden_state_quantized->type = input_weights->type;
+    hidden_state_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(hidden_state_quantized->dims,
+                             hidden_state->dims)) {
+      TfLiteIntArray* hidden_state_quantized_size =
+          TfLiteIntArrayCopy(hidden_state->dims);
+      TF_LITE_ENSURE_OK(context,
+                        context->ResizeTensor(context, hidden_state_quantized,
+                                              hidden_state_quantized_size));
+    }
+    node->temporaries->data[2] = op_data->scratch_tensor_index + 2;
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/2,
+                                                &scaling_factors));
+    scaling_factors->type = kTfLiteFloat32;
+    scaling_factors->allocation_type = kTfLiteArenaRw;
+    int scaling_dims[1] = {batch_size};
+    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
+      TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1);
+      scaling_factors_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors,
+                                                       scaling_factors_size));
+    }
+    node->temporaries->data[3] = op_data->scratch_tensor_index + 3;
+    TfLiteTensor* accum_scratch;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/3, &accum_scratch));
+    accum_scratch->type = kTfLiteInt32;
+    accum_scratch->allocation_type = kTfLiteArenaRw;
+    int accum_scratch_dims[2] = {num_units, batch_size};
+    if (!TfLiteIntArrayEqualsArray(accum_scratch->dims, 2,
+                                   accum_scratch_dims)) {
+      TfLiteIntArray* accum_scratch_size = TfLiteIntArrayCreate(2);
+      accum_scratch_size->data[0] = accum_scratch_dims[0];
+      accum_scratch_size->data[1] = accum_scratch_dims[1];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, accum_scratch,
+                                                       accum_scratch_size));
+    }
+    node->temporaries->data[4] = op_data->scratch_tensor_index + 4;
+    TfLiteTensor* zero_points;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/4, &zero_points));
+    zero_points->type = kTfLiteInt32;
+    zero_points->allocation_type = kTfLiteArenaRw;
+    int zero_points_dims[1] = {batch_size};
+    if (!TfLiteIntArrayEqualsArray(zero_points->dims, 1, zero_points_dims)) {
+      TfLiteIntArray* zero_points_size = TfLiteIntArrayCreate(1);
+      zero_points_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, zero_points,
+                                                       zero_points_size));
+    }
+    node->temporaries->data[5] = op_data->scratch_tensor_index + 5;
+    TfLiteTensor* row_sums;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, /*index=*/5, &row_sums));
+    row_sums->type = kTfLiteInt32;
+    row_sums->name = "Rnn_row_sums";
+    row_sums->allocation_type = kTfLiteArenaRwPersistent;
+    int row_sums_dims[2] = {2, num_units};
+    if (!TfLiteIntArrayEqualsArray(row_sums->dims, 2, row_sums_dims)) {
+      TfLiteIntArray* row_sums_size = TfLiteIntArrayCreate(2);
+      row_sums_size->data[0] = row_sums_dims[0];
+      row_sums_size->data[1] = row_sums_dims[1];
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, row_sums, row_sums_size));
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalFloat(const TfLiteTensor* input,
+                       const TfLiteTensor* input_weights,
+                       const TfLiteTensor* recurrent_weights,
+                       const TfLiteTensor* bias, const TfLiteRNNParams* params,
+                       TfLiteTensor* hidden_state, TfLiteTensor* output) {
+  const int batch_size = input->dims->data[0];
+  const int num_units = input_weights->dims->data[0];
+  const int input_size = input->dims->data[1];
+  const int output_batch_leading_dim =
+      output->dims->data[output->dims->size - 1];
+
+  // Initialize the pointer to hidden state.
+  float* hidden_state_ptr_batch = GetTensorData<float>(hidden_state);
+  // Initialize the pointer to input and output.
+  const float* input_ptr_batch = GetTensorData<float>(input);
+  float* output_ptr_batch = GetTensorData<float>(output);
+  // Initialize input_weights, recurrent_weights and bias.
+  const float* input_weights_ptr = GetTensorData<float>(input_weights);
+  const float* recurrent_weights_ptr = GetTensorData<float>(recurrent_weights);
+  const float* bias_ptr = GetTensorData<float>(bias);
+
+  kernel_utils::RnnBatchStep(
+      input_ptr_batch, input_weights_ptr, recurrent_weights_ptr, bias_ptr,
+      input_size, num_units, batch_size, output_batch_leading_dim,
+      params->activation, hidden_state_ptr_batch, output_ptr_batch);
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalHybrid(const TfLiteTensor* input,
+                        const TfLiteTensor* input_weights,
+                        const TfLiteTensor* recurrent_weights,
+                        const TfLiteTensor* bias, const TfLiteRNNParams* params,
+                        TfLiteTensor* input_scratch,
+                        TfLiteTensor* hidden_state_scratch,
+                        TfLiteTensor* scaling_factors,
+                        TfLiteTensor* hidden_state, TfLiteTensor* output,
+                        TfLiteTensor* zero_points, TfLiteTensor* accum_scratch,
+                        TfLiteTensor* row_sums, bool* compute_row_sums) {
+  const int batch_size = input->dims->data[0];
+  const int num_units = input_weights->dims->data[0];
+  const int input_size = input->dims->data[1];
+  const int output_batch_leading_dim =
+      output->dims->data[output->dims->size - 1];
+
+  // Initialize the pointer to hidden state.
+  float* hidden_state_ptr_batch = GetTensorData<float>(hidden_state);
+  // Initialize the pointer to input and output.
+  const float* input_ptr_batch = GetTensorData<float>(input);
+  float* output_ptr_batch = GetTensorData<float>(output);
+  // Initialize input_weights, recurrent_weights and bias.
+  const int8_t* input_weights_ptr = GetTensorData<int8_t>(input_weights);
+  const int8_t* recurrent_weights_ptr =
+      GetTensorData<int8_t>(recurrent_weights);
+  const float* bias_ptr = GetTensorData<float>(bias);
+  // Get the scale of the quantized weights.
+  float input_weights_scale = input_weights->params.scale;
+  float recurrent_weights_scale = recurrent_weights->params.scale;
+  // Initialize temporary storage for quantized values.
+  int8_t* quantized_input_ptr = GetTensorData<int8_t>(input_scratch);
+  int8_t* quantized_hidden_state_ptr =
+      GetTensorData<int8_t>(hidden_state_scratch);
+  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
+  int32_t* accum_scratch_ptr = GetTensorData<int32_t>(accum_scratch);
+  int32_t* zero_points_ptr = nullptr;
+  int32_t* row_sums_ptr = nullptr;
+  if (params->asymmetric_quantize_inputs) {
+    zero_points_ptr = GetTensorData<int32_t>(zero_points);
+    row_sums_ptr = GetTensorData<int32_t>(row_sums);
+  }
+  kernel_utils::RnnBatchStep(
+      input_ptr_batch, input_weights_ptr, input_weights_scale,
+      recurrent_weights_ptr, recurrent_weights_scale, bias_ptr, input_size,
+      num_units, batch_size, output_batch_leading_dim, params->activation,
+      quantized_input_ptr, quantized_hidden_state_ptr, scaling_factors_ptr,
+      hidden_state_ptr_batch, output_ptr_batch,
+      params->asymmetric_quantize_inputs, zero_points_ptr, accum_scratch_ptr,
+      row_sums_ptr, compute_row_sums);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params = reinterpret_cast<TfLiteRNNParams*>(node->builtin_data);
+  auto* op_data = reinterpret_cast<OpData*>(node->user_data);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* input_weights;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node, kWeightsTensor, &input_weights));
+  const TfLiteTensor* recurrent_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node, kRecurrentWeightsTensor, &recurrent_weights));
+  const TfLiteTensor* bias;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBiasTensor, &bias));
+  TfLiteTensor* hidden_state =
+      GetVariableInput(context, node, kHiddenStateTensor);
+  TF_LITE_ENSURE(context, hidden_state != nullptr);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // We already checked that weight types are consistent, so branch on one.
+  switch (input_weights->type) {
+    case kTfLiteFloat32:
+      return EvalFloat(input, input_weights, recurrent_weights, bias, params,
+                       hidden_state, output);
+    case kTfLiteUInt8:
+    case kTfLiteInt8: {
+      // TODO(mirkov): implement eval with quantized inputs as well.
+      TfLiteTensor* input_quantized;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, 0, &input_quantized));
+      TfLiteTensor* hidden_state_quantized;
+      TF_LITE_ENSURE_OK(
+          context, GetTemporarySafe(context, node, 1, &hidden_state_quantized));
+      TfLiteTensor* scaling_factors;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, 2, &scaling_factors));
+      TfLiteTensor* accum_scratch;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, 3, &accum_scratch));
+      TfLiteTensor* zero_points;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, 4, &zero_points));
+      TfLiteTensor* row_sums;
+      TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, 5, &row_sums));
+      return EvalHybrid(input, input_weights, recurrent_weights, bias, params,
+                        input_quantized, hidden_state_quantized,
+                        scaling_factors, hidden_state, output, zero_points,
+                        accum_scratch, row_sums, &op_data->compute_row_sums);
+    }
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %s not currently supported.",
+                         TfLiteTypeGetName(input_weights->type));
+      return kTfLiteError;
+  }
+}
+
+}  // namespace rnn
+
+TfLiteRegistration* Register_RNN() {
+  static TfLiteRegistration r = {rnn::Init, rnn::Free, rnn::Prepare, rnn::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/basic_rnn_test.cc

@@ -0,0 +1,348 @@
+/* Copyright 2017 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.
+==============================================================================*/
+// Unit test for TFLite RNN op.
+
+#include <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+static float rnn_input[] = {
+    0.23689353,   0.285385,     0.037029743, -0.19858193,  -0.27569133,
+    0.43773448,   0.60379338,   0.35562468,  -0.69424844,  -0.93421471,
+    -0.87287879,  0.37144363,   -0.62476718, 0.23791671,   0.40060222,
+    0.1356622,    -0.99774903,  -0.98858172, -0.38952237,  -0.47685933,
+    0.31073618,   0.71511042,   -0.63767755, -0.31729108,  0.33468103,
+    0.75801885,   0.30660987,   -0.37354088, 0.77002847,   -0.62747043,
+    -0.68572164,  0.0069220066, 0.65791464,  0.35130811,   0.80834007,
+    -0.61777675,  -0.21095741,  0.41213346,  0.73784804,   0.094794154,
+    0.47791874,   0.86496925,   -0.53376222, 0.85315156,   0.10288584,
+    0.86684,      -0.011186242, 0.10513687,  0.87825835,   0.59929144,
+    0.62827742,   0.18899453,   0.31440187,  0.99059987,   0.87170351,
+    -0.35091716,  0.74861872,   0.17831337,  0.2755419,    0.51864719,
+    0.55084288,   0.58982027,   -0.47443086, 0.20875752,   -0.058871567,
+    -0.66609079,  0.59098077,   0.73017097,  0.74604273,   0.32882881,
+    -0.17503482,  0.22396147,   0.19379807,  0.29120302,   0.077113032,
+    -0.70331609,  0.15804303,   -0.93407321, 0.40182066,   0.036301374,
+    0.66521823,   0.0300982,    -0.7747041,  -0.02038002,  0.020698071,
+    -0.90300065,  0.62870288,   -0.23068321, 0.27531278,   -0.095755219,
+    -0.712036,    -0.17384434,  -0.50593495, -0.18646687,  -0.96508682,
+    0.43519354,   0.14744234,   0.62589407,  0.1653645,    -0.10651493,
+    -0.045277178, 0.99032974,   -0.88255352, -0.85147917,  0.28153265,
+    0.19455957,   -0.55479527,  -0.56042433, 0.26048636,   0.84702539,
+    0.47587705,   -0.074295521, -0.12287641, 0.70117295,   0.90532446,
+    0.89782166,   0.79817224,   0.53402734,  -0.33286154,  0.073485017,
+    -0.56172788,  -0.044897556, 0.89964068,  -0.067662835, 0.76863563,
+    0.93455386,   -0.6324693,   -0.083922029};
+
+static float rnn_golden_output[] = {
+    0.496726,   0,          0.965996,  0,         0.0584254, 0,
+    0,          0.12315,    0,         0,         0.612266,  0.456601,
+    0,          0.52286,    1.16099,   0.0291232,
+
+    0,          0,          0.524901,  0,         0,         0,
+    0,          1.02116,    0,         1.35762,   0,         0.356909,
+    0.436415,   0.0355727,  0,         0,
+
+    0,          0,          0,         0.262335,  0,         0,
+    0,          1.33992,    0,         2.9739,    0,         0,
+    1.31914,    2.66147,    0,         0,
+
+    0.942568,   0,          0,         0,         0.025507,  0,
+    0,          0,          0.321429,  0.569141,  1.25274,   1.57719,
+    0.8158,     1.21805,    0.586239,  0.25427,
+
+    1.04436,    0,          0.630725,  0,         0.133801,  0.210693,
+    0.363026,   0,          0.533426,  0,         1.25926,   0.722707,
+    0,          1.22031,    1.30117,   0.495867,
+
+    0.222187,   0,          0.72725,   0,         0.767003,  0,
+    0,          0.147835,   0,         0,         0,         0.608758,
+    0.469394,   0.00720298, 0.927537,  0,
+
+    0.856974,   0.424257,   0,         0,         0.937329,  0,
+    0,          0,          0.476425,  0,         0.566017,  0.418462,
+    0.141911,   0.996214,   1.13063,   0,
+
+    0.967899,   0,          0,         0,         0.0831304, 0,
+    0,          1.00378,    0,         0,         0,         1.44818,
+    1.01768,    0.943891,   0.502745,  0,
+
+    0.940135,   0,          0,         0,         0,         0,
+    0,          2.13243,    0,         0.71208,   0.123918,  1.53907,
+    1.30225,    1.59644,    0.70222,   0,
+
+    0.804329,   0,          0.430576,  0,         0.505872,  0.509603,
+    0.343448,   0,          0.107756,  0.614544,  1.44549,   1.52311,
+    0.0454298,  0.300267,   0.562784,  0.395095,
+
+    0.228154,   0,          0.675323,  0,         1.70536,   0.766217,
+    0,          0,          0,         0.735363,  0.0759267, 1.91017,
+    0.941888,   0,          0,         0,
+
+    0,          0,          1.5909,    0,         0,         0,
+    0,          0.5755,     0,         0.184687,  0,         1.56296,
+    0.625285,   0,          0,         0,
+
+    0,          0,          0.0857888, 0,         0,         0,
+    0,          0.488383,   0.252786,  0,         0,         0,
+    1.02817,    1.85665,    0,         0,
+
+    0.00981836, 0,          1.06371,   0,         0,         0,
+    0,          0,          0,         0.290445,  0.316406,  0,
+    0.304161,   1.25079,    0.0707152, 0,
+
+    0.986264,   0.309201,   0,         0,         0,         0,
+    0,          1.64896,    0.346248,  0,         0.918175,  0.78884,
+    0.524981,   1.92076,    2.07013,   0.333244,
+
+    0.415153,   0.210318,   0,         0,         0,         0,
+    0,          2.02616,    0,         0.728256,  0.84183,   0.0907453,
+    0.628881,   3.58099,    1.49974,   0};
+
+static std::initializer_list<float> rnn_weights = {
+    0.461459,    0.153381,   0.529743,    -0.00371218, 0.676267,   -0.211346,
+    0.317493,    0.969689,   -0.343251,   0.186423,    0.398151,   0.152399,
+    0.448504,    0.317662,   0.523556,    -0.323514,   0.480877,   0.333113,
+    -0.757714,   -0.674487,  -0.643585,   0.217766,    -0.0251462, 0.79512,
+    -0.595574,   -0.422444,  0.371572,    -0.452178,   -0.556069,  -0.482188,
+    -0.685456,   -0.727851,  0.841829,    0.551535,    -0.232336,  0.729158,
+    -0.00294906, -0.69754,   0.766073,    -0.178424,   0.369513,   -0.423241,
+    0.548547,    -0.0152023, -0.757482,   -0.85491,    0.251331,   -0.989183,
+    0.306261,    -0.340716,  0.886103,    -0.0726757,  -0.723523,  -0.784303,
+    0.0354295,   0.566564,   -0.485469,   -0.620498,   0.832546,   0.697884,
+    -0.279115,   0.294415,   -0.584313,   0.548772,    0.0648819,  0.968726,
+    0.723834,    -0.0080452, -0.350386,   -0.272803,   0.115121,   -0.412644,
+    -0.824713,   -0.992843,  -0.592904,   -0.417893,   0.863791,   -0.423461,
+    -0.147601,   -0.770664,  -0.479006,   0.654782,    0.587314,   -0.639158,
+    0.816969,    -0.337228,  0.659878,    0.73107,     0.754768,   -0.337042,
+    0.0960841,   0.368357,   0.244191,    -0.817703,   -0.211223,  0.442012,
+    0.37225,     -0.623598,  -0.405423,   0.455101,    0.673656,   -0.145345,
+    -0.511346,   -0.901675,  -0.81252,    -0.127006,   0.809865,   -0.721884,
+    0.636255,    0.868989,   -0.347973,   -0.10179,    -0.777449,  0.917274,
+    0.819286,    0.206218,   -0.00785118, 0.167141,    0.45872,    0.972934,
+    -0.276798,   0.837861,   0.747958,    -0.0151566,  -0.330057,  -0.469077,
+    0.277308,    0.415818};
+
+static std::initializer_list<float> rnn_recurrent_weights = {
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+    0.1};
+
+static std::initializer_list<float> rnn_bias = {
+    0.065691948, -0.69055247, 0.1107955,  -0.97084129, -0.23957068, -0.23566568,
+    -0.389184,   0.47481549,  -0.4791103, 0.29931796,  0.10463274,  0.83918178,
+    0.37197268,  0.61957061,  0.3956964,  -0.37609905};
+
+class RNNOpModel : public SingleOpModel {
+ public:
+  RNNOpModel(int batches, int units, int size,
+             const TensorType& weights = TensorType_FLOAT32,
+             const TensorType& recurrent_weights = TensorType_FLOAT32,
+             bool asymmetric_quantize_inputs = false)
+      : batches_(batches), units_(units), input_size_(size) {
+    input_ = AddInput(TensorType_FLOAT32);
+    weights_ = AddInput(weights);
+    recurrent_weights_ = AddInput(recurrent_weights);
+    bias_ = AddInput(TensorType_FLOAT32);
+    hidden_state_ = AddVariableInput(TensorType_FLOAT32);
+    output_ = AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_RNN, BuiltinOptions_RNNOptions,
+                 CreateRNNOptions(builder_, ActivationFunctionType_RELU,
+                                  asymmetric_quantize_inputs)
+                     .Union());
+    BuildInterpreter({{batches_, input_size_},  // input tensor
+                      {units_, input_size_},    // weights tensor
+                      {units_, units_},         // recurrent weights tensor
+                      {units_},                 // bias tensor
+                      {batches_, units_}});     // hidden state tensor
+  }
+
+  void SetBias(std::initializer_list<float> f) { PopulateTensor(bias_, f); }
+
+  void SetWeights(std::initializer_list<float> f) {
+    PopulateTensor(weights_, f);
+  }
+
+  void SetRecurrentWeights(std::initializer_list<float> f) {
+    PopulateTensor(recurrent_weights_, f);
+  }
+
+  void SetInput(std::initializer_list<float> data) {
+    PopulateTensor(input_, data);
+  }
+
+  void SetInput(int offset, float* begin, float* end) {
+    PopulateTensor(input_, offset, begin, end);
+  }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+
+  int input_size() { return input_size_; }
+  int num_units() { return units_; }
+  int num_batches() { return batches_; }
+
+ protected:
+  int input_;
+  int weights_;
+  int recurrent_weights_;
+  int bias_;
+  int hidden_state_;
+  int output_;
+
+  int batches_;
+  int units_;
+  int input_size_;
+};
+
+// The hybrid model has quantized weights and recurrent_weights.
+class HybridRNNOpModel : public RNNOpModel {
+ public:
+  HybridRNNOpModel(int batches, int units, int size, TensorType tensor_type,
+                   bool asymmetric_quantize_inputs)
+      : RNNOpModel(batches, units, size, tensor_type, tensor_type,
+                   asymmetric_quantize_inputs) {
+    tensor_type_ = tensor_type;
+  }
+
+  TensorType tensor_type_;
+
+  void SetWeights(int weights_idx, const std::vector<float>& f) {
+    if (tensor_type_ == TensorType_UINT8) {
+      SymmetricQuantizeAndPopulate(weights_idx, f);
+    } else {
+      SignedSymmetricQuantizeAndPopulate(weights_idx, f);
+    }
+  }
+
+  void SetWeights(std::initializer_list<float> f) { SetWeights(weights_, f); }
+
+  void SetRecurrentWeights(std::initializer_list<float> f) {
+    SetWeights(recurrent_weights_, f);
+  }
+};
+
+TEST(RnnOpTest, BlackBoxTest) {
+  RNNOpModel rnn(2, 16, 8);
+  rnn.SetWeights(rnn_weights);
+  rnn.SetBias(rnn_bias);
+  rnn.SetRecurrentWeights(rnn_recurrent_weights);
+
+  const int input_sequence_size = sizeof(rnn_input) / sizeof(float) /
+                                  (rnn.input_size() * rnn.num_batches());
+
+  for (int i = 0; i < input_sequence_size; i++) {
+    float* batch_start = rnn_input + i * rnn.input_size();
+    float* batch_end = batch_start + rnn.input_size();
+    rnn.SetInput(0, batch_start, batch_end);
+    rnn.SetInput(rnn.input_size(), batch_start, batch_end);
+
+    ASSERT_EQ(rnn.Invoke(), kTfLiteOk);
+
+    float* golden_start = rnn_golden_output + i * rnn.num_units();
+    float* golden_end = golden_start + rnn.num_units();
+    std::vector<float> expected;
+    expected.insert(expected.end(), golden_start, golden_end);
+    expected.insert(expected.end(), golden_start, golden_end);
+
+    EXPECT_THAT(rnn.GetOutput(), ElementsAreArray(ArrayFloatNear(expected)));
+  }
+}
+
+class HybridRnnOpTest : public ::testing::TestWithParam<bool> {};
+
+TEST_P(HybridRnnOpTest, BlackBoxTestUint8) {
+  HybridRNNOpModel rnn(2, 16, 8, TensorType_UINT8, GetParam());
+  rnn.SetWeights(rnn_weights);
+  rnn.SetBias(rnn_bias);
+  rnn.SetRecurrentWeights(rnn_recurrent_weights);
+
+  const int input_sequence_size = sizeof(rnn_input) / sizeof(float) /
+                                  (rnn.input_size() * rnn.num_batches());
+
+  for (int i = 0; i < input_sequence_size; i++) {
+    float* batch_start = rnn_input + i * rnn.input_size();
+    float* batch_end = batch_start + rnn.input_size();
+    rnn.SetInput(0, batch_start, batch_end);
+    rnn.SetInput(rnn.input_size(), batch_start, batch_end);
+
+    ASSERT_EQ(rnn.Invoke(), kTfLiteOk);
+
+    float* golden_start = rnn_golden_output + i * rnn.num_units();
+    float* golden_end = golden_start + rnn.num_units();
+    std::vector<float> expected;
+    expected.insert(expected.end(), golden_start, golden_end);
+    expected.insert(expected.end(), golden_start, golden_end);
+
+    EXPECT_THAT(rnn.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                     expected, /*max_abs_err=*/0.0104)));
+  }
+}
+
+TEST_P(HybridRnnOpTest, BlackBoxTestInt8) {
+  HybridRNNOpModel rnn(2, 16, 8, TensorType_INT8, GetParam());
+  rnn.SetWeights(rnn_weights);
+  rnn.SetBias(rnn_bias);
+  rnn.SetRecurrentWeights(rnn_recurrent_weights);
+
+  const int input_sequence_size = sizeof(rnn_input) / sizeof(float) /
+                                  (rnn.input_size() * rnn.num_batches());
+
+  for (int i = 0; i < input_sequence_size; i++) {
+    float* batch_start = rnn_input + i * rnn.input_size();
+    float* batch_end = batch_start + rnn.input_size();
+    rnn.SetInput(0, batch_start, batch_end);
+    rnn.SetInput(rnn.input_size(), batch_start, batch_end);
+
+    ASSERT_EQ(rnn.Invoke(), kTfLiteOk);
+
+    float* golden_start = rnn_golden_output + i * rnn.num_units();
+    float* golden_end = golden_start + rnn.num_units();
+    std::vector<float> expected;
+    expected.insert(expected.end(), golden_start, golden_end);
+    expected.insert(expected.end(), golden_start, golden_end);
+
+    EXPECT_THAT(rnn.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                     expected, /*max_abs_err=*/0.0104)));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(HybridRnnOpTest, HybridRnnOpTest,
+                         ::testing::ValuesIn({false, true}));
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/batch_matmul.cc

@@ -0,0 +1,843 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/reference/batch_matmul.h"
+
+#include <stddef.h>
+#include <string.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <limits>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace batch_matmul {
+
+static const int kInputLHSTensor = 0;
+static const int kInputRHSTensor = 1;
+static const int kOutputTensor = 0;
+
+static const int kNumTempTensorsForAdjoints = 2;
+static const int kNumTempTensorsForHybrid = 5;
+
+struct OpData {
+  // The scaling factor from input to output (aka the 'real multiplier') can
+  // be represented as a fixed point multiplier plus a left shift.
+  int32_t output_multiplier;
+  int output_shift;
+  // The range of the fused activation layer. For example for kNone and
+  // uint8_t these would be 0 and 255.
+  int32_t output_activation_min;
+  int32_t output_activation_max;
+  // The index of the temporary tensors where we store transposed LHS/RHS.
+  int scratch_tensor_index;
+  bool rhs_transposed;
+  bool compute_row_sums = false;
+};
+
+struct OpContext {
+  OpContext(TfLiteContext* context, TfLiteNode* node) {
+    params = reinterpret_cast<TfLiteBatchMatMulParams*>(node->builtin_data);
+    lhs = GetInput(context, node, kInputLHSTensor);
+    rhs = GetInput(context, node, kInputRHSTensor);
+    output = GetOutput(context, node, 0);
+  }
+  TfLiteBatchMatMulParams* params;
+  const TfLiteTensor* lhs;
+  const TfLiteTensor* rhs;
+  TfLiteTensor* output;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  // If the RHS is constant, we only transpose once.
+  op_data->rhs_transposed = false;
+  // Creates the temp tensors to store the transposed LHS and/or RHS, and
+  // extra buffers for the quantized case.
+  op_data->scratch_tensor_index = -1;
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete static_cast<OpData*>(buffer);
+}
+
+TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
+                                const RuntimeShape& extended_lhs_shape,
+                                const RuntimeShape& extended_rhs_shape,
+                                bool adj_x, bool adj_y, int output_rank,
+                                TfLiteTensor* output) {
+  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(output_rank);
+  // Fill in any broadcast dimensions.
+  for (int i = 0; i < output_rank - 2; ++i) {
+    const int lhs_dim = extended_lhs_shape.Dims(i);
+    const int rhs_dim = extended_rhs_shape.Dims(i);
+    int broadcast_dim = lhs_dim;
+    if ((lhs_dim != rhs_dim) && (lhs_dim == 1)) {
+      broadcast_dim = rhs_dim;
+    }
+    output_shape->data[i] = broadcast_dim;
+  }
+  // Fill in the matmul dimensions.
+  int lhs_rows_index = adj_x ? output_rank - 1 : output_rank - 2;
+  int rhs_cols_index = adj_y ? output_rank - 2 : output_rank - 1;
+
+  output_shape->data[output_rank - 2] = extended_lhs_shape.Dims(lhs_rows_index);
+  output_shape->data[output_rank - 1] = extended_rhs_shape.Dims(rhs_cols_index);
+  TfLiteStatus stat = context->ResizeTensor(context, output, output_shape);
+  return stat;
+}
+
+// Initializes temp tensors to store transposed operands.
+TfLiteStatus InitializeTemporaries(TfLiteContext* context, TfLiteNode* node,
+                                   OpContext* op_context) {
+  // Create temporary tensors to hold transposed LHS/RHS.
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  const TfLiteTensor* lhs = op_context->lhs;
+  const TfLiteTensor* rhs = op_context->rhs;
+  TfLiteIntArrayFree(node->temporaries);
+  // For "hybrid" quantization, we impose the constraint that the LHS
+  // is float (typically an activation from a prior layer) and the RHS
+  // is quantized int8.
+  bool is_hybrid =
+      (op_context->lhs->type == kTfLiteFloat32 && rhs->type == kTfLiteInt8);
+  if (is_hybrid) {
+    node->temporaries = TfLiteIntArrayCreate(kNumTempTensorsForAdjoints +
+                                             kNumTempTensorsForHybrid);
+  } else {
+    node->temporaries = TfLiteIntArrayCreate(kNumTempTensorsForAdjoints);
+  }
+
+  const int lhs_rank = NumDimensions(lhs);
+  const int rhs_rank = NumDimensions(rhs);
+  const int batch_size = op_context->params->adj_x
+                             ? lhs->dims->data[lhs_rank - 1]
+                             : lhs->dims->data[lhs_rank - 2];
+  const int num_units = op_context->params->adj_y
+                            ? rhs->dims->data[rhs_rank - 2]
+                            : rhs->dims->data[rhs_rank - 1];
+
+  // Temp tensor for Transposed LHS;
+  {
+    node->temporaries->data[0] = op_data->scratch_tensor_index;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/0, &scratch_buffer));
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(lhs_rank);
+    for (int i = 0; i < lhs_rank - 2; ++i) {
+      scratch_buffer_size->data[i] = lhs->dims->data[i];
+    }
+    // Swap last two dimensions.
+    scratch_buffer_size->data[lhs_rank - 2] = lhs->dims->data[lhs_rank - 1];
+    scratch_buffer_size->data[lhs_rank - 1] = lhs->dims->data[lhs_rank - 2];
+
+    scratch_buffer->type = op_context->lhs->type;
+    scratch_buffer->allocation_type = kTfLiteArenaRw;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // We need a temp buffer for the RHS if we need to transpose the RHS. We
+  // transpose by default, so that the two inputs (LHS and RHS) are in a proper
+  // layout for our fast matrix multiplication routines. If the transpose flag
+  // is set by the caller, the data is already in the desired layout.
+  {
+    node->temporaries->data[1] = op_data->scratch_tensor_index + 1;
+    TfLiteTensor* scratch_buffer;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/1, &scratch_buffer));
+    scratch_buffer->name = "BatchMatMul_scratch_buffer";
+    const TfLiteTensor* rhs = op_context->rhs;
+    int rhs_rank = NumDimensions(rhs);
+    TfLiteIntArray* scratch_buffer_size = TfLiteIntArrayCreate(rhs_rank);
+    for (int i = 0; i < rhs_rank - 2; ++i) {
+      scratch_buffer_size->data[i] = rhs->dims->data[i];
+    }
+    // Swap last two dimensions.
+    scratch_buffer_size->data[rhs_rank - 2] = rhs->dims->data[rhs_rank - 1];
+    scratch_buffer_size->data[rhs_rank - 1] = rhs->dims->data[rhs_rank - 2];
+
+    if (IsConstantTensor(op_context->rhs)) {
+      scratch_buffer->allocation_type = kTfLiteArenaRwPersistent;
+    } else {
+      scratch_buffer->allocation_type = kTfLiteArenaRw;
+    }
+    scratch_buffer->type = op_context->rhs->type;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_buffer,
+                                                     scratch_buffer_size));
+  }
+
+  // If we have to perform on-the-fly quantization (with quantized weights and
+  // float inputs) first we need to quantize the inputs. Allocate temporary
+  // buffer to store the intermediate quantized values, the batch scaling
+  // factors, the input offsets, and the sums of the rows for each weights
+  // matrix.
+  // RHS = weights, LHS = inputs
+  if (is_hybrid) {
+    // Calculate the total number of LHS batches.
+    int num_batches = 1;
+    for (int i = 0; i < lhs_rank - 2; ++i) {
+      num_batches *= lhs->dims->data[i];
+    }
+    int num_weights_matrices = 1;
+    for (int i = 0; i < rhs_rank - 2; ++i) {
+      num_weights_matrices *= rhs->dims->data[i];
+    }
+    op_data->compute_row_sums = true;
+    node->temporaries->data[2] = op_data->scratch_tensor_index + 2;
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/2,
+                                                &input_quantized));
+    input_quantized->type = op_context->rhs->type;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+
+    TfLiteIntArray* input_quantized_size =
+        TfLiteIntArrayCopy(op_context->lhs->dims);
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                     input_quantized_size));
+
+    node->temporaries->data[3] = op_data->scratch_tensor_index + 3;
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/3,
+                                                &scaling_factors));
+    scaling_factors->type = kTfLiteFloat32;
+    scaling_factors->allocation_type = kTfLiteArenaRw;
+    // Total size of scaling factors is batch size * number of total batches
+    int scaling_dims[1] = {num_batches * batch_size};
+    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
+      TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1);
+      scaling_factors_size->data[0] = scaling_dims[0];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors,
+                                                       scaling_factors_size));
+    }
+
+    node->temporaries->data[4] = op_data->scratch_tensor_index + 4;
+    TfLiteTensor* accum_scratch;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/4, &accum_scratch));
+    accum_scratch->type = kTfLiteInt32;
+    accum_scratch->allocation_type = kTfLiteArenaRw;
+    int accum_scratch_dims[2] = {num_units, batch_size};
+    if (!TfLiteIntArrayEqualsArray(accum_scratch->dims, 2,
+                                   accum_scratch_dims)) {
+      TfLiteIntArray* accum_size = TfLiteIntArrayCreate(2);
+      accum_size->data[0] = num_units;
+      accum_size->data[1] = batch_size;
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, accum_scratch, accum_size));
+    }
+
+    node->temporaries->data[5] = op_data->scratch_tensor_index + 5;
+    TfLiteTensor* input_offsets;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/5, &input_offsets));
+    input_offsets->type = kTfLiteInt32;
+    input_offsets->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqualsArray(input_offsets->dims, 1, scaling_dims)) {
+      TfLiteIntArray* input_offsets_size = TfLiteIntArrayCreate(1);
+      input_offsets_size->data[0] = num_batches * batch_size;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_offsets,
+                                                       input_offsets_size));
+    }
+    node->temporaries->data[6] = op_data->scratch_tensor_index + 6;
+    TfLiteTensor* row_sums;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, /*index=*/6, &row_sums));
+    row_sums->type = kTfLiteInt32;
+    row_sums->allocation_type = kTfLiteArenaRwPersistent;
+    int row_sums_dims[1] = {num_weights_matrices * num_units};
+    if (!TfLiteIntArrayEqualsArray(row_sums->dims, 1, row_sums_dims)) {
+      TfLiteIntArray* row_sums_size = TfLiteIntArrayCreate(1);
+      row_sums_size->data[0] = row_sums_dims[0];
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, row_sums, row_sums_size));
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  if (op_data->scratch_tensor_index == -1) {
+    context->AddTensors(context,
+                        kNumTempTensorsForAdjoints + kNumTempTensorsForHybrid,
+                        &op_data->scratch_tensor_index);
+  }
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  OpContext op_context(context, node);
+  TF_LITE_ENSURE_OK(context, InitializeTemporaries(context, node, &op_context));
+
+  bool adj_x = op_context.params->adj_x;
+  bool adj_y = op_context.params->adj_y;
+
+  const TfLiteTensor* lhs_data;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputLHSTensor, &lhs_data));
+  const TfLiteTensor* rhs_data;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputRHSTensor, &rhs_data));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // Note that quantized inference requires that all tensors have their
+  // parameters set. This is usually done during quantized training.
+  if ((lhs_data->type == kTfLiteInt8 || lhs_data->type == kTfLiteInt16) &&
+      output->type != kTfLiteInt32) {
+    double real_multiplier = 0.0;
+    TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler(
+        context, lhs_data, rhs_data, output, &real_multiplier));
+    int exponent;
+    QuantizeMultiplier(real_multiplier, &op_data->output_multiplier, &exponent);
+    op_data->output_shift = exponent;
+    // BatchMatMul has no fused activation functions. Therefore, set
+    // output activation min and max to min and max of int8_t or int16_t
+    // type.
+    if (lhs_data->type == kTfLiteInt8) {
+      op_data->output_activation_min = std::numeric_limits<int8_t>::min();
+      op_data->output_activation_max = std::numeric_limits<int8_t>::max();
+    } else {
+      op_data->output_activation_min = std::numeric_limits<int16_t>::min();
+      op_data->output_activation_max = std::numeric_limits<int16_t>::max();
+    }
+  }
+
+  if (lhs_data->type == kTfLiteInt16) {
+    TF_LITE_ENSURE_EQ(context, lhs_data->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, rhs_data->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
+  }
+
+  TF_LITE_ENSURE(context, lhs_data->type == kTfLiteFloat32 ||
+                              lhs_data->type == kTfLiteInt8 ||
+                              lhs_data->type == kTfLiteInt16);
+  TF_LITE_ENSURE(context, rhs_data->type == kTfLiteFloat32 ||
+                              rhs_data->type == kTfLiteInt8 ||
+                              rhs_data->type == kTfLiteInt16);
+  // Either we have a hybrid quantization with a float32 and an int8 input,
+  // otherwise both inputs should be of the same type.
+  TF_LITE_ENSURE(
+      context,
+      (lhs_data->type == kTfLiteFloat32 && rhs_data->type == kTfLiteInt8) ||
+          lhs_data->type == rhs_data->type ||
+          (lhs_data->type == kTfLiteInt16 && rhs_data->type == kTfLiteInt8));
+  // Support dimensions between 2 and 5, inclusive.
+  TF_LITE_ENSURE(context, NumDimensions(lhs_data) >= 2);
+  TF_LITE_ENSURE(context, NumDimensions(lhs_data) <= 5);
+  TF_LITE_ENSURE(context, NumDimensions(rhs_data) >= 2);
+  TF_LITE_ENSURE(context, NumDimensions(rhs_data) <= 5);
+
+  const int lhs_rank = NumDimensions(lhs_data);
+  const int rhs_rank = NumDimensions(rhs_data);
+  const int output_rank = std::max(lhs_rank, rhs_rank);
+  const RuntimeShape extended_lhs_shape =
+      RuntimeShape::ExtendedShape(output_rank, GetTensorShape(lhs_data));
+  const RuntimeShape extended_rhs_shape =
+      RuntimeShape::ExtendedShape(output_rank, GetTensorShape(rhs_data));
+
+  // Ensure any batch dimensions obey broadcasting rules.
+  for (int i = 0; i < output_rank - 2; ++i) {
+    const int lhs_dim = extended_lhs_shape.Dims(i);
+    const int rhs_dim = extended_rhs_shape.Dims(i);
+    if (lhs_dim != rhs_dim) {
+      if (lhs_dim != 1) {
+        TF_LITE_ENSURE_EQ(context, rhs_dim, 1);
+      }
+    }
+  }
+  // Ensure other dimensions work for matrix multiplication.
+  int accum_dim_lhs = adj_x ? extended_lhs_shape.Dims(output_rank - 2)
+                            : extended_lhs_shape.Dims(output_rank - 1);
+  int accum_dim_rhs = adj_y ? extended_rhs_shape.Dims(output_rank - 1)
+                            : extended_rhs_shape.Dims(output_rank - 2);
+
+  TF_LITE_ENSURE_EQ(context, accum_dim_lhs, accum_dim_rhs);
+  TfLiteStatus status =
+      ResizeOutputTensor(context, extended_lhs_shape, extended_rhs_shape, adj_x,
+                         adj_y, output_rank, output);
+  return status;
+}
+
+template <typename scalar>
+void TransposeRowsColumnsImpl(const TfLiteTensor* tensor_in,
+                              const scalar* input, TfLiteTensor* tensor_out,
+                              scalar* output) {
+  RuntimeShape transposed_shape(GetTensorShape(tensor_in));
+  RuntimeShape shape(GetTensorShape(tensor_in));
+  TransposeParams params;
+  int rank = NumDimensions(tensor_in);
+  params.perm_count = rank;
+  for (int i = 0; i < rank - 2; ++i) {
+    params.perm[i] = i;
+  }
+  // Transpose the last two dimensions.
+  params.perm[rank - 2] = rank - 1;
+  params.perm[rank - 1] = rank - 2;
+  transposed_shape.SetDim(rank - 1, shape.Dims(rank - 2));
+  transposed_shape.SetDim(rank - 2, shape.Dims(rank - 1));
+  optimized_ops::Transpose(params, shape, input, transposed_shape, output);
+}
+
+TfLiteStatus TransposeRowsColumns(TfLiteContext* context,
+                                  const TfLiteTensor* tensor_in,
+                                  TfLiteTensor* tensor_out) {
+  if (tensor_in->type == kTfLiteFloat32) {
+    TransposeRowsColumnsImpl<float>(tensor_in, GetTensorData<float>(tensor_in),
+                                    tensor_out,
+                                    GetTensorData<float>(tensor_out));
+    return kTfLiteOk;
+  } else if (tensor_in->type == kTfLiteInt8) {
+    TransposeRowsColumnsImpl<int8_t>(
+        tensor_in, GetTensorData<int8_t>(tensor_in), tensor_out,
+        GetTensorData<int8_t>(tensor_out));
+    return kTfLiteOk;
+  } else if (tensor_in->type == kTfLiteInt16) {
+    TransposeRowsColumnsImpl<int16_t>(
+        tensor_in, GetTensorData<int16_t>(tensor_in), tensor_out,
+        GetTensorData<int16_t>(tensor_out));
+    return kTfLiteOk;
+  } else {
+    TF_LITE_KERNEL_LOG(
+        context, "Can only transpose tensors with float, int8 or int16 type.");
+    return kTfLiteError;
+  }
+}
+
+RuntimeShape SwapRowColumnDims(const RuntimeShape& shape) {
+  RuntimeShape swapped_shape(shape);
+  const int32_t dims = shape.DimensionsCount();
+  swapped_shape.SetDim(dims - 2, shape.Dims(dims - 1));
+  swapped_shape.SetDim(dims - 1, shape.Dims(dims - 2));
+  return swapped_shape;
+}
+
+TfLiteStatus VerifyPerChannelQuantization(TfLiteContext* context,
+                                          const TfLiteTensor* tensor) {
+  TF_LITE_ENSURE_EQ(context, tensor->quantization.type,
+                    kTfLiteAffineQuantization);
+  const auto* affine_quantization =
+      reinterpret_cast<TfLiteAffineQuantization*>(tensor->quantization.params);
+  TF_LITE_ENSURE(context, affine_quantization);
+  TF_LITE_ENSURE(context, affine_quantization->scale);
+  return affine_quantization->scale->size > 1 ? kTfLiteOk : kTfLiteError;
+}
+
+TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, OpData* data,
+                        const RuntimeShape& input_shape,
+                        const TfLiteTensor* input,
+                        const RuntimeShape& filter_shape,
+                        const TfLiteTensor* filter,
+                        TfLiteTensor* input_quantized,
+                        TfLiteTensor* scaling_factors,
+                        TfLiteTensor* accum_scratch, TfLiteTensor* row_sums,
+                        TfLiteTensor* input_offsets, TfLiteTensor* output) {
+  const auto* params =
+      reinterpret_cast<TfLiteBatchMatMulParams*>(node->builtin_data);
+  const int32_t num_input_dims = input_shape.DimensionsCount();
+
+  // Input row/cols have been swapped at this point, so dims are
+  // {input_size, num_batches}
+  const int input_size = input_shape.Dims(num_input_dims - 2);
+  const int batch_size = input_shape.Dims(num_input_dims - 1);
+
+  int num_batches_to_quantize = batch_size;
+  for (int i = 0; i < input_shape.DimensionsCount() - 2; ++i) {
+    num_batches_to_quantize *= input_shape.Dims(i);
+  }
+  // Quantize input from float to uint8 + quantization params (scaling factor).
+  const int scaling_factor_size = GetTensorShape(scaling_factors).FlatSize();
+  TF_LITE_ENSURE(context, scaling_factor_size >= num_batches_to_quantize);
+  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
+  int32_t* input_offset_ptr = nullptr;
+  int32_t* row_sums_ptr = nullptr;
+  input_offset_ptr = GetTensorData<int32_t>(input_offsets);
+  row_sums_ptr = GetTensorData<int32_t>(row_sums);
+  if (!params->asymmetric_quantize_inputs) {
+    memset(input_offset_ptr, 0, input_offsets->bytes);
+  }
+  int8_t* quant_data = GetTensorData<int8_t>(input_quantized);
+  const int8_t* filter_data = GetTensorData<int8_t>(filter);
+  const float* input_ptr = GetTensorData<float>(input);
+  // Quantize each batch independently.
+  tensor_utils::BatchQuantizeFloats(input_ptr, num_batches_to_quantize,
+                                    input_size, quant_data, scaling_factors_ptr,
+                                    input_offset_ptr,
+                                    params->asymmetric_quantize_inputs);
+  float* per_channel_scale_ptr = nullptr;
+  if (VerifyPerChannelQuantization(context, filter) == kTfLiteOk) {
+    //  Per channel quantization.
+    const auto* affine_quantization =
+        reinterpret_cast<TfLiteAffineQuantization*>(
+            filter->quantization.params);
+    TF_LITE_ENSURE_EQ(
+        context, affine_quantization->scale->size,
+        filter->dims->data[affine_quantization->quantized_dimension]);
+    per_channel_scale_ptr = affine_quantization->scale->data;
+  } else {
+    // Per tensor quantization.
+    for (int b = 0; b < num_batches_to_quantize; ++b) {
+      // Incorporate scaling of the filter
+      scaling_factors_ptr[b] *= filter->params.scale;
+    }
+  }
+
+  RuntimeShape output_shape = GetTensorShape(output);
+  int output_size = 1;
+  for (int i = 0; i < output_shape.DimensionsCount(); ++i) {
+    output_size *= output_shape.Dims(i);
+  }
+  std::fill_n(GetTensorData<float>(output), output_size, 0.0f);
+  reference_ops::BatchMatMul(filter_shape, filter_data, input_shape, quant_data,
+                             scaling_factors_ptr, input_offset_ptr,
+                             row_sums_ptr, GetTensorShape(output),
+                             GetTensorData<float>(output),
+                             &(data->compute_row_sums), per_channel_scale_ptr);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalInt8Int8(TfLiteContext* context, const OpData* data,
+                          const RuntimeShape& lhs_shape,
+                          const TfLiteTensor* lhs,
+                          const RuntimeShape& rhs_shape,
+                          const TfLiteTensor* rhs,
+                          const RuntimeShape& output_shape,
+                          TfLiteTensor* output) {
+  // Reuse params struct from FullyConnected Op.
+  FullyConnectedParams op_params;
+  int32_t input_offset = -lhs->params.zero_point;
+  int32_t filter_offset = -rhs->params.zero_point;
+  int32_t output_offset = output->params.zero_point;
+  op_params.input_offset = input_offset;
+  op_params.weights_offset = filter_offset;
+  op_params.output_offset = output_offset;
+  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_shift = data->output_shift;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+  op_params.lhs_cacheable = IsConstantTensor(lhs);
+  op_params.rhs_cacheable = IsConstantTensor(rhs);
+
+  reference_ops::BatchMatMul<int8_t, int32_t>(
+      op_params, rhs_shape, GetTensorData<int8_t>(rhs), lhs_shape,
+      GetTensorData<int8_t>(lhs), GetTensorShape(output),
+      GetTensorData<int8_t>(output));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalInt8Int32(TfLiteContext* context, const OpData* data,
+                           const RuntimeShape& lhs_shape,
+                           const TfLiteTensor* lhs,
+                           const RuntimeShape& rhs_shape,
+                           const TfLiteTensor* rhs,
+                           const RuntimeShape& output_shape,
+                           TfLiteTensor* output) {
+  // Set BatchMatMul lhs param to rhs(filter) and rhs param to lhs(input). For
+  // the reason, see comment of Eval() function.
+  reference_ops::BatchMatMul<int8, int8, int32>(
+      rhs_shape, GetTensorData<int8>(rhs), lhs_shape, GetTensorData<int8>(lhs),
+      GetTensorShape(output), GetTensorData<int32>(output));
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalInt16(TfLiteContext* context, const OpData* data,
+                       const RuntimeShape& lhs_shape, const TfLiteTensor* lhs,
+                       const RuntimeShape& rhs_shape, const TfLiteTensor* rhs,
+                       const RuntimeShape& output_shape, TfLiteTensor* output) {
+  // Reuse params struct from FullyConnected Op.
+  FullyConnectedParams op_params;
+  int32_t input_offset = -lhs->params.zero_point;
+  int32_t filter_offset = -rhs->params.zero_point;
+  int32_t output_offset = output->params.zero_point;
+  op_params.input_offset = input_offset;
+  op_params.weights_offset = filter_offset;
+  op_params.output_offset = output_offset;
+  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_shift = data->output_shift;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+
+  reference_ops::BatchMatMul<int16_t, int64_t>(
+      op_params, rhs_shape, GetTensorData<int16_t>(rhs), lhs_shape,
+      GetTensorData<int16_t>(lhs), GetTensorShape(output),
+      GetTensorData<int16_t>(output));
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalInt16Int8(TfLiteContext* context, const OpData* data,
+                           const RuntimeShape& lhs_shape,
+                           const TfLiteTensor* lhs,
+                           const RuntimeShape& rhs_shape,
+                           const TfLiteTensor* rhs,
+                           const RuntimeShape& output_shape,
+                           TfLiteTensor* output) {
+  // Reuse params struct from FullyConnected Op.
+  FullyConnectedParams op_params;
+  op_params.input_offset = -lhs->params.zero_point;
+  op_params.weights_offset = -rhs->params.zero_point;
+  op_params.output_offset = output->params.zero_point;
+  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_shift = data->output_shift;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+
+  reference_ops::BatchMatMul<int8_t, int64_t, int16_t, int16_t>(
+      op_params, rhs_shape, GetTensorData<int8_t>(rhs), lhs_shape,
+      GetTensorData<int16_t>(lhs), GetTensorShape(output),
+      GetTensorData<int16_t>(output));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
+                           OpData* data, const RuntimeShape& lhs_shape,
+                           const TfLiteTensor* lhs,
+                           const RuntimeShape& rhs_shape,
+                           const TfLiteTensor* rhs, TfLiteTensor* output) {
+  if (lhs->type == kTfLiteFloat32 && rhs->type == kTfLiteInt8) {
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/2,
+                                                &input_quantized));
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, /*index=*/3,
+                                                &scaling_factors));
+    TfLiteTensor* accum_scratch;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/4, &accum_scratch));
+    TfLiteTensor* input_offsets;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, /*index=*/5, &input_offsets));
+    TfLiteTensor* row_sums;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, /*index=*/6, &row_sums));
+    return EvalHybrid(context, node, data, lhs_shape, lhs, rhs_shape, rhs,
+                      input_quantized, scaling_factors, accum_scratch, row_sums,
+                      input_offsets, output);
+  } else if (lhs->type == kTfLiteInt8 && rhs->type == kTfLiteInt8) {
+    if (output->type == kTfLiteInt8) {
+      return EvalInt8Int8(context, data, lhs_shape, lhs, rhs_shape, rhs,
+                          GetTensorShape(output), output);
+    } else {
+      return EvalInt8Int32(context, data, lhs_shape, lhs, rhs_shape, rhs,
+                           GetTensorShape(output), output);
+    }
+  } else if (lhs->type == kTfLiteInt16 && rhs->type == kTfLiteInt16) {
+    return EvalInt16(context, data, lhs_shape, lhs, rhs_shape, rhs,
+                     GetTensorShape(output), output);
+  } else if (lhs->type == kTfLiteInt16 && rhs->type == kTfLiteInt8) {
+    return EvalInt16Int8(context, data, lhs_shape, lhs, rhs_shape, rhs,
+                         GetTensorShape(output), output);
+  } else {
+    TF_LITE_KERNEL_LOG(
+        context,
+        "Currently only hybrid, int8 and int16 quantization are supported.\n");
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteTensor* GetTempRhs(TfLiteContext* context, TfLiteNode* node,
+                         const TfLiteTensor* rhs) {
+  TfLiteTensor* transposed_rhs = GetTemporary(context, node, 1);
+  if (transposed_rhs == nullptr) {
+    return nullptr;
+  }
+
+  TfLiteIntArrayFree(transposed_rhs->dims);
+  transposed_rhs->dims = TfLiteIntArrayCopy(rhs->dims);
+  std::swap(transposed_rhs->dims->data[transposed_rhs->dims->size - 1],
+            transposed_rhs->dims->data[transposed_rhs->dims->size - 2]);
+  if (rhs->type == kTfLiteInt8 || rhs->type == kTfLiteInt16) {
+    // Get the quantization params from the RHS tensor.
+    transposed_rhs->params.scale = rhs->params.scale;
+    transposed_rhs->params.zero_point = rhs->params.zero_point;
+    if (rhs->quantization.type == kTfLiteAffineQuantization) {
+      transposed_rhs->quantization.type = rhs->quantization.type;
+      if (transposed_rhs->quantization.params) {
+        auto* transposed_rhs_affine_quantization =
+            reinterpret_cast<TfLiteAffineQuantization*>(
+                transposed_rhs->quantization.params);
+        TfLiteIntArrayFree(transposed_rhs_affine_quantization->zero_point);
+        TfLiteFloatArrayFree(transposed_rhs_affine_quantization->scale);
+        free(transposed_rhs->quantization.params);
+      }
+      transposed_rhs->quantization.params =
+          malloc(sizeof(TfLiteAffineQuantization));
+      const auto* rhs_affine_quantization =
+          reinterpret_cast<TfLiteAffineQuantization*>(rhs->quantization.params);
+      auto* transposed_rhs_affine_quantization =
+          reinterpret_cast<TfLiteAffineQuantization*>(
+              transposed_rhs->quantization.params);
+      int quantized_dimension = rhs_affine_quantization->quantized_dimension;
+      if (quantized_dimension == rhs->dims->size - 1) {
+        quantized_dimension = rhs->dims->size - 2;
+      } else if (quantized_dimension == rhs->dims->size - 2) {
+        quantized_dimension = rhs->dims->size - 1;
+      }
+      transposed_rhs_affine_quantization->quantized_dimension =
+          quantized_dimension;
+      transposed_rhs_affine_quantization->zero_point =
+          TfLiteIntArrayCopy(rhs_affine_quantization->zero_point);
+      transposed_rhs_affine_quantization->scale =
+          TfLiteFloatArrayCopy(rhs_affine_quantization->scale);
+    }
+  }
+  return transposed_rhs;
+}
+
+TfLiteTensor* GetTempLhs(TfLiteContext* context, TfLiteNode* node,
+                         const TfLiteTensor* lhs) {
+  TfLiteTensor* transposed_lhs = GetTemporary(context, node, 0);
+  if (transposed_lhs == nullptr) {
+    return nullptr;
+  }
+
+  if (lhs->type == kTfLiteInt8 || lhs->type == kTfLiteInt16) {
+    // Get the quantization params from the LHS tensor.
+    transposed_lhs->params.scale = lhs->params.scale;
+    transposed_lhs->params.zero_point = lhs->params.zero_point;
+  }
+  return transposed_lhs;
+}
+
+// Perform a batch matrix multiply on
+// LHS <..., A, B>  X  RHS<..., B, C>
+// where the leading dimensions of LHS and RHS obey broadcasting rules
+// (this Op will apply broadcasting rules).
+// We assume that LHS and RHS are both row oriented (adjacent values in memory
+// are in the same row) and will output in the same memory layout. However,
+// our fast GEMM libraries assume RCC layout (LHS row oriented,
+// RHS column oriented, output column oriented). Therefore, we perform
+// RHS <..., C, B> X LHS <..., B, A>
+// where output is a C X A column-oriented, which is equivalent to
+// A X C row-oriented.
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  OpContext op_context(context, node);
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  const TfLiteTensor* lhs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputLHSTensor, &lhs));
+  const TfLiteTensor* rhs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputRHSTensor, &rhs));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  RuntimeShape orig_lhs_shape = GetTensorShape(lhs);
+  RuntimeShape orig_rhs_shape = GetTensorShape(rhs);
+
+  bool adj_y = op_context.params->adj_y;
+  bool adj_x = op_context.params->adj_x;
+
+  int32_t rhs_dims_count = orig_rhs_shape.DimensionsCount();
+  int32_t lhs_dims_count = orig_lhs_shape.DimensionsCount();
+  // Compress ops where rhs shape is [..., 1, X, Y] and lhs shape is
+  // [..., Q, R, S] which is equivalent to rhs: [..., X, Y] and
+  // lhs: [..., Q * R, S].
+  // This compression is only valid when adj_x is false. If we squash the
+  // batch dimension 'Q' into the row dimension 'R' and then transpose the
+  // resulting [..., Q * R, S] matrix, we will incorrectly scramble the data
+  // across batches.
+  if (!adj_x && rhs_dims_count > 2 && lhs_dims_count > 2) {
+    int rhs_one = orig_rhs_shape.DimsData()[rhs_dims_count - 3];
+    if (rhs_one == 1) {
+      int32_t* lhs_dims = orig_lhs_shape.DimsData();
+      int32_t* rhs_dims = orig_rhs_shape.DimsData();
+      RuntimeShape tmp_l(lhs_dims_count - 1, lhs_dims);
+      tmp_l.SetDim(lhs_dims_count - 3,
+                   lhs_dims[lhs_dims_count - 3] * lhs_dims[lhs_dims_count - 2]);
+      tmp_l.SetDim(lhs_dims_count - 2, lhs_dims[lhs_dims_count - 1]);
+      orig_lhs_shape.ReplaceWith(tmp_l.DimensionsCount(), tmp_l.DimsData());
+      RuntimeShape tmp_r(rhs_dims_count - 1, orig_rhs_shape.DimsData());
+      tmp_r.SetDim(rhs_dims_count - 3, rhs_dims[rhs_dims_count - 2]);
+      tmp_r.SetDim(rhs_dims_count - 2, rhs_dims[rhs_dims_count - 1]);
+      orig_rhs_shape.ReplaceWith(tmp_r.DimensionsCount(), tmp_r.DimsData());
+    }
+  }
+  rhs_dims_count = orig_rhs_shape.DimensionsCount();
+  lhs_dims_count = orig_lhs_shape.DimensionsCount();
+  const TfLiteTensor* lhs_tensor = adj_x ? GetTempLhs(context, node, lhs) : lhs;
+  const TfLiteTensor* rhs_tensor = adj_y ? rhs : GetTempRhs(context, node, rhs);
+  if (!adj_y) {
+    // TODO(b/154760341) Constant tensors should already be transposed, but
+    // we transpose once if necessary for now.
+    if (!(IsConstantTensor(rhs) && op_data->rhs_transposed)) {
+      TransposeRowsColumns(context, rhs, GetTemporary(context, node, 1));
+      op_data->rhs_transposed = true;
+    }
+  }
+  if (adj_x) {
+    TransposeRowsColumns(context, lhs, GetTemporary(context, node, 0));
+  }
+  RuntimeShape rhs_shape =
+      adj_y ? orig_rhs_shape : SwapRowColumnDims(orig_rhs_shape);
+  RuntimeShape lhs_shape =
+      adj_x ? orig_lhs_shape : SwapRowColumnDims(orig_lhs_shape);
+
+  switch (rhs->type) {
+    case kTfLiteFloat32:
+      // Note we pass RHS args first, LHS args second. See note above.
+      reference_ops::BatchMatMul(rhs_shape, GetTensorData<float>(rhs_tensor),
+                                 lhs_shape, GetTensorData<float>(lhs_tensor),
+                                 GetTensorShape(output),
+                                 GetTensorData<float>(output));
+      break;
+    case kTfLiteInt8:
+    case kTfLiteInt16:
+      EvalQuantized(context, node, op_data, lhs_shape, lhs_tensor, rhs_shape,
+                    rhs_tensor, output);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Currently BatchMatMul doesn't support type: %s",
+                         TfLiteTypeGetName(lhs->type));
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace batch_matmul
+
+TfLiteRegistration* Register_BATCH_MATMUL_REF() {
+  static TfLiteRegistration r = {batch_matmul::Init, batch_matmul::Free,
+                                 batch_matmul::Prepare, batch_matmul::Eval};
+  return &r;
+}
+
+TfLiteRegistration* Register_BATCH_MATMUL() {
+  return Register_BATCH_MATMUL_REF();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/batch_to_space_nd.cc

@@ -0,0 +1,208 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace batch_to_space_nd {
+
+// This file has two implementations of BatchToSpaceND.
+enum KernelType {
+  kReference,
+  kGenericOptimized,
+};
+
+struct BatchToSpaceNDContext {
+  BatchToSpaceNDContext(TfLiteContext* context, TfLiteNode* node) {
+    input = GetInput(context, node, 0);
+    block_shape = GetInput(context, node, 1);
+    crops = GetInput(context, node, 2);
+    output = GetOutput(context, node, 0);
+  }
+  const TfLiteTensor* input;
+  const TfLiteTensor* block_shape;
+  const TfLiteTensor* crops;
+  TfLiteTensor* output;
+};
+
+// Currently, only 3D NHC or 4D NHWC input/output op_context are supported.
+// In case of 3D input,it will be converted to 4D by adding W=1 to be NH1C.
+// The 4D array need to have exactly 2 spatial dimensions.
+// TODO(ycling): Support arbitrary dimension in BatchToSpaceND.
+const int kInputMinDimensionNum = 3;
+const int kInputMaxDimensionNum = 4;
+
+TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
+                                BatchToSpaceNDContext* op_context) {
+  TfLiteIntArray* input_size = op_context->input->dims;
+  const int* block_shape = GetTensorData<int32>(op_context->block_shape);
+  const int* crops = GetTensorData<int32>(op_context->crops);
+
+  int spatial_dims_num = input_size->size - 2;
+  // Block_shape should be a 1D tensor with dimension [spatial_dims_num].
+  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->block_shape), 1);
+  TF_LITE_ENSURE_EQ(context, op_context->block_shape->dims->data[0],
+                    spatial_dims_num);
+  // Crops should be a 2D tensor with dimension [spatial_dims_num, 2].
+  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->crops), 2);
+  TF_LITE_ENSURE_EQ(context, op_context->crops->dims->data[0],
+                    spatial_dims_num);
+  TF_LITE_ENSURE_EQ(context, op_context->crops->dims->data[1], 2);
+
+  for (int i = 0; i < spatial_dims_num * 2; ++i) {
+    TF_LITE_ENSURE(context, crops[i] >= 0);
+  }
+
+  TfLiteIntArray* output_size = TfLiteIntArrayCopy(input_size);
+  int output_batch_size = input_size->data[0];
+  for (int dim = 0; dim < spatial_dims_num; ++dim) {
+    // Number of batch must be multiple of (block_shape[dim]).
+    TF_LITE_ENSURE(context, block_shape[dim] != 0);
+    TF_LITE_ENSURE_EQ(context, output_batch_size % block_shape[dim], 0);
+    output_batch_size = output_batch_size / block_shape[dim];
+    output_size->data[dim + 1] = input_size->data[dim + 1] * block_shape[dim] -
+                                 crops[dim * 2] - crops[dim * 2 + 1];
+  }
+
+  output_size->data[0] = output_batch_size;
+  output_size->data[input_size->size - 1] =
+      input_size->data[input_size->size - 1];
+
+  return context->ResizeTensor(context, op_context->output, output_size);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  BatchToSpaceNDContext op_context(context, node);
+  TF_LITE_ENSURE(context,
+                 NumDimensions(op_context.input) >= kInputMinDimensionNum);
+  TF_LITE_ENSURE(context,
+                 NumDimensions(op_context.input) <= kInputMaxDimensionNum);
+  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+
+  if (op_context.input->type == kTfLiteUInt8 ||
+      op_context.input->type == kTfLiteInt8 ||
+      op_context.input->type == kTfLiteInt16) {
+    TF_LITE_ENSURE_EQ(context, op_context.input->params.scale,
+                      op_context.output->params.scale);
+    TF_LITE_ENSURE_EQ(context, op_context.input->params.zero_point,
+                      op_context.output->params.zero_point);
+  }
+
+  if (op_context.input->type == kTfLiteInt16) {
+    TF_LITE_ENSURE_EQ(context, op_context.input->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, op_context.output->params.zero_point, 0);
+  }
+
+  if (!IsConstantOrPersistentTensor(op_context.block_shape) ||
+      !IsConstantOrPersistentTensor(op_context.crops)) {
+    SetTensorToDynamic(op_context.output);
+    return kTfLiteOk;
+  }
+  return ResizeOutputTensor(context, &op_context);
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  BatchToSpaceNDContext op_context(context, node);
+
+  // Resize the output tensor if the output tensor is dynamic.
+  if (IsDynamicTensor(op_context.output)) {
+    TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
+  }
+
+#define TF_LITE_BATCH_TO_SPACE_ND(type, scalar)                        \
+  type::BatchToSpaceND(GetTensorShape(op_context.input),               \
+                       GetTensorData<scalar>(op_context.input),        \
+                       GetTensorShape(op_context.block_shape),         \
+                       GetTensorData<int32_t>(op_context.block_shape), \
+                       GetTensorShape(op_context.crops),               \
+                       GetTensorData<int32_t>(op_context.crops),       \
+                       GetTensorShape(op_context.output),              \
+                       GetTensorData<scalar>(op_context.output))
+  // Already know in/out types are same.
+  switch (TfLiteTypeGetSizeBits(op_context.input->type)) {
+    case 8:
+      if (kernel_type == kReference) {
+        TF_LITE_BATCH_TO_SPACE_ND(reference_ops, uint8_t);
+      } else {
+        TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, uint8_t);
+      }
+      break;
+    case 16:
+      if (kernel_type == kReference) {
+        TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int16_t);
+      } else {
+        TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int16_t);
+      }
+      break;
+    case 32:
+      if (kernel_type == kReference) {
+        TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int32_t);
+      } else {
+        TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int32_t);
+      }
+      break;
+    case 64:
+      if (kernel_type == kReference) {
+        TF_LITE_BATCH_TO_SPACE_ND(reference_ops, int64_t);
+      } else {
+        TF_LITE_BATCH_TO_SPACE_ND(optimized_ops, int64_t);
+      }
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Type %d is currently not supported by BatchToSpace.",
+                         op_context.input->type);
+      return kTfLiteError;
+  }
+#undef TF_LITE_BATCH_TO_SPACE_ND
+  return kTfLiteOk;
+}
+
+}  // namespace batch_to_space_nd
+
+TfLiteRegistration* Register_BATCH_TO_SPACE_ND_REF() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, batch_to_space_nd::Prepare,
+      batch_to_space_nd::Eval<batch_to_space_nd::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_BATCH_TO_SPACE_ND_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, batch_to_space_nd::Prepare,
+      batch_to_space_nd::Eval<batch_to_space_nd::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_BATCH_TO_SPACE_ND() {
+  return Register_BATCH_TO_SPACE_ND_GENERIC_OPT();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/batch_to_space_nd_test.cc

@@ -0,0 +1,356 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include <initializer_list>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class BatchToSpaceNDOpModel : public SingleOpModel {
+ public:
+  template <typename T>
+  void SetInput(std::initializer_list<T> data) {
+    PopulateTensor<T>(input_, data);
+  }
+
+  void SetBlockShape(std::initializer_list<int> data) {
+    PopulateTensor<int>(block_shape_, data);
+  }
+
+  void SetCrops(std::initializer_list<int> data) {
+    PopulateTensor<int>(crops_, data);
+  }
+
+  int input_tensor_id() { return input_; }
+
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+
+  template <typename T>
+  std::vector<float> GetDequantizedOutput() {
+    return Dequantize<T>(ExtractVector<T>(output_), GetScale(output_),
+                         GetZeroPoint(output_));
+  }
+
+  int32_t GetOutputSize() { return GetTensorSize(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  int input_;
+  int block_shape_;
+  int crops_;
+  int output_;
+};
+
+// Tests case where block_shape and crops are const tensors.
+//
+// Example usage is as follows:
+//    BatchToSpaceNDOpConstModel m(input_shape, block_shape, crops);
+//    m.SetInput(input_data);
+//    m.Invoke();
+class BatchToSpaceNDOpConstModel : public BatchToSpaceNDOpModel {
+ public:
+  BatchToSpaceNDOpConstModel(const TensorData& input,
+                             std::initializer_list<int> block_shape,
+                             std::initializer_list<int> crops,
+                             const TensorData& output) {
+    int spatial_dims = static_cast<int>(block_shape.size());
+    input_ = AddInput(input);
+    block_shape_ = AddConstInput(TensorType_INT32, block_shape, {spatial_dims});
+    crops_ = AddConstInput(TensorType_INT32, crops, {spatial_dims, 2});
+    output_ = AddOutput(output);
+
+    SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND,
+                 BuiltinOptions_BatchToSpaceNDOptions,
+                 CreateBatchToSpaceNDOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+};
+
+// Tests case where block_shape and crops are non-const tensors.
+//
+// Example usage is as follows:
+//    BatchToSpaceNDOpDynamicModel m(input_shape);
+//    m.SetInput(input_data);
+//    m.SetBlockShape(block_shape);
+//    m.SetPaddings(crops);
+//    m.Invoke();
+class BatchToSpaceNDOpDynamicModel : public BatchToSpaceNDOpModel {
+ public:
+  BatchToSpaceNDOpDynamicModel(const TensorData& input,
+                               const TensorData& output) {
+    input_ = AddInput(input);
+    block_shape_ = AddInput(TensorType_INT32);
+    crops_ = AddInput(TensorType_INT32);
+    output_ = AddOutput(output);
+
+    int spatial_dims = static_cast<int>(GetShape(input_).size()) - 2;
+    SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND,
+                 BuiltinOptions_BatchToSpaceNDOptions,
+                 CreateBatchToSpaceNDOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_), {spatial_dims}, {spatial_dims, 2}});
+  }
+};
+
+template <typename integer_type>
+float GetTolerance(float min, float max) {
+  float kQuantizedStep =
+      (max - min) / (std::numeric_limits<integer_type>::max() -
+                     std::numeric_limits<integer_type>::min());
+  return kQuantizedStep;
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleConstTest) {
+  BatchToSpaceNDOpConstModel m({TensorType_FLOAT32, {4, 2, 2, 1}}, {2, 2},
+                               {0, 0, 0, 0}, {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray(
+                  {1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleConstTestInt8) {
+  BatchToSpaceNDOpConstModel m({TensorType_INT8, {4, 2, 2, 1}}, {2, 2},
+                               {0, 0, 0, 0}, {TensorType_INT8});
+  m.SetInput<int8_t>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+  EXPECT_THAT(m.GetOutput<int8_t>(),
+              ElementsAreArray(
+                  {1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, BatchOneConstTest) {
+  BatchToSpaceNDOpConstModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, {1, 1},
+                               {0, 0, 0, 0}, {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2, 2, 1}));
+  EXPECT_THAT(m.GetOutput<float>(), ElementsAreArray({1, 2, 3, 4}));
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleConstTestInt8EmptyOutput) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    // NNAPI doesn't currently support non-zero crop values.
+    return;
+  }
+
+  BatchToSpaceNDOpConstModel m({TensorType_INT8, {4, 2, 2, 1}}, {2, 2},
+                               {0, 0, 2, 2}, {TensorType_INT8});
+  m.SetInput<int8_t>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 0, 1}));
+  EXPECT_THAT(m.GetOutputSize(), 0);
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleDynamicTest) {
+  BatchToSpaceNDOpDynamicModel m({TensorType_FLOAT32, {4, 2, 2, 1}},
+                                 {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2, 2});
+  m.SetCrops({0, 0, 0, 0});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray(
+                  {1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleDynamicTestInt8) {
+  BatchToSpaceNDOpDynamicModel m({TensorType_INT8, {4, 2, 2, 1}},
+                                 {TensorType_INT8});
+  m.SetInput<int8_t>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2, 2});
+  m.SetCrops({0, 0, 0, 0});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1}));
+  EXPECT_THAT(m.GetOutput<int8_t>(),
+              ElementsAreArray(
+                  {1, 5, 2, 6, 9, 13, 10, 14, 3, 7, 4, 8, 11, 15, 12, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, InvalidCropsDynamicTest) {
+  BatchToSpaceNDOpDynamicModel m({TensorType_FLOAT32, {4, 2, 2, 1}},
+                                 {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2, 2});
+  m.SetCrops({0, 0, -1, 0});
+  ASSERT_NE(m.Invoke(), kTfLiteOk) << "crops.i. >= 0 was not true.";
+}
+
+TEST(BatchToSpaceNDOpTest, SimpleDynamicTestInt8EmptyOutput) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    // NNAPI doesn't currently support non-zero crop values.
+    return;
+  }
+
+  BatchToSpaceNDOpDynamicModel m({TensorType_INT8, {4, 2, 2, 1}},
+                                 {TensorType_INT8});
+  m.SetInput<int8_t>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2, 2});
+  m.SetCrops({2, 2, 0, 0});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 0, 4, 1}));
+  EXPECT_THAT(m.GetOutput<int8_t>(), ::testing::IsEmpty());
+}
+
+#if GTEST_HAS_DEATH_TEST
+TEST(BatchToSpaceNDOpTest, InvalidShapeTest) {
+  EXPECT_DEATH(
+      BatchToSpaceNDOpConstModel({TensorType_FLOAT32, {3, 2, 2, 1}}, {2, 2},
+                                 {0, 0, 0, 0}, {TensorType_FLOAT32}),
+      "Cannot allocate tensors");
+}
+
+TEST(BatchToSpaceNDOpTest, InvalidCropsConstTest) {
+  EXPECT_DEATH(
+      BatchToSpaceNDOpConstModel({TensorType_FLOAT32, {3, 2, 2, 1}}, {2, 2},
+                                 {0, 0, 0, -1}, {TensorType_FLOAT32}),
+      "crops.i. >= 0 was not true.");
+}
+#endif
+
+TEST(BatchToSpaceNDOpTest, Simple3DConstTest) {
+  BatchToSpaceNDOpConstModel m({TensorType_FLOAT32, {4, 4, 1}}, {2}, {0, 0},
+                               {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 8, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray(
+                  {1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15, 8, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DConstTestWithCrops) {
+  BatchToSpaceNDOpConstModel m({TensorType_FLOAT32, {4, 4, 1}}, {2}, {1, 1},
+                               {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 6, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray({9, 2, 10, 3, 11, 4, 13, 6, 14, 7, 15, 8}));
+}
+
+template <typename integer_dtype>
+void Simple3DConstTestWithCropsQuant() {
+  const float kMin = -1;
+  const float kMax =
+      std::numeric_limits<integer_dtype>::max() /
+      static_cast<float>(std::numeric_limits<integer_dtype>::max() + 1);
+  float kQuantizedTolerance = GetTolerance<integer_dtype>(-16.0f, 16.0f);
+  BatchToSpaceNDOpConstModel m(
+      {GetTensorType<integer_dtype>(), {4, 4, 1}, 16.0f * kMin, 16.0f * kMax},
+      {2}, {1, 1},
+      {GetTensorType<integer_dtype>(), {}, 16.0f * kMin, 16.0f * kMax});
+  m.QuantizeAndPopulate<integer_dtype>(
+      m.input_tensor_id(),
+      {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 6, 1}));
+  EXPECT_THAT(
+      m.GetDequantizedOutput<integer_dtype>(),
+      ElementsAreArray(ArrayFloatNear({9, 2, 10, 3, 11, 4, 13, 6, 14, 7, 15, 8},
+                                      kQuantizedTolerance)));
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DConstTestWithCropsUINT8) {
+  Simple3DConstTestWithCropsQuant<uint8_t>();
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DConstTestWithCropsINT8) {
+  Simple3DConstTestWithCropsQuant<int8_t>();
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DConstTestWithCropsINT16) {
+  Simple3DConstTestWithCropsQuant<int16_t>();
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DDynamicTest) {
+  BatchToSpaceNDOpDynamicModel m({TensorType_FLOAT32, {4, 4, 1}},
+                                 {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2});
+  m.SetCrops({0, 0});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 8, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray(
+                  {1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15, 8, 16}));
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DDynamicTestWithCrops) {
+  BatchToSpaceNDOpDynamicModel m({TensorType_FLOAT32, {4, 4, 1}},
+                                 {TensorType_FLOAT32});
+  m.SetInput<float>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2});
+  m.SetCrops({1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 6, 1}));
+  EXPECT_THAT(m.GetOutput<float>(),
+              ElementsAreArray({9, 2, 10, 3, 11, 4, 13, 6, 14, 7, 15, 8}));
+}
+
+template <typename integer_dtype>
+void Simple3DDynamicTestWithCropsQuant() {
+  const float kMin = -1;
+  const float kMax =
+      std::numeric_limits<integer_dtype>::max() /
+      static_cast<float>(std::numeric_limits<integer_dtype>::max() + 1);
+  float kQuantizedTolerance = GetTolerance<integer_dtype>(-16.0, 16.0);
+  BatchToSpaceNDOpDynamicModel m(
+      {GetTensorType<integer_dtype>(), {4, 4, 1}, 16.0f * kMin, 16.0f * kMax},
+      {GetTensorType<integer_dtype>(), {}, 16.0f * kMin, 16.0f * kMax});
+  m.QuantizeAndPopulate<integer_dtype>(
+      m.input_tensor_id(),
+      {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetBlockShape({2});
+  m.SetCrops({1, 1});
+  m.Invoke();
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 6, 1}));
+  EXPECT_THAT(
+      m.GetDequantizedOutput<integer_dtype>(),
+      ElementsAreArray(ArrayFloatNear({9, 2, 10, 3, 11, 4, 13, 6, 14, 7, 15, 8},
+                                      kQuantizedTolerance)));
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DDynamicTestWithCropsQuantUINT8) {
+  Simple3DDynamicTestWithCropsQuant<uint8_t>();
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DDynamicTestWithCropsQuantINT8) {
+  Simple3DDynamicTestWithCropsQuant<int8_t>();
+}
+
+TEST(BatchToSpaceNDOpTest, Simple3DDynamicTestWithCropsQuantINT16) {
+  Simple3DDynamicTestWithCropsQuant<int16_t>();
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/bidirectional_sequence_rnn.cc

@@ -0,0 +1,851 @@
+/* Copyright 2017 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 <algorithm>
+#include <cstddef>
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace bidirectional_sequence_rnn {
+
+namespace {
+
+struct OpData {
+  int scratch_tensor_index;
+  bool fw_compute_row_sums = false;
+  bool bw_compute_row_sums = false;
+};
+
+}  // namespace
+
+// LINT.IfChange
+
+constexpr int kInputTensor = 0;
+// Forward and backward cell tensors.
+constexpr int kFwWeightsTensor = 1;
+constexpr int kFwRecurrentWeightsTensor = 2;
+constexpr int kFwBiasTensor = 3;
+constexpr int kFwHiddenStateTensor = 4;
+constexpr int kBwWeightsTensor = 5;
+constexpr int kBwRecurrentWeightsTensor = 6;
+constexpr int kBwBiasTensor = 7;
+constexpr int kBwHiddenStateTensor = 8;
+// Used as auxiliary input and weights when stacking for
+// tf.contrib.rnn.stack_bidirectional_rnn case (with cross links); Used as input
+// to the backward cell when stacking for tf.nn.static_bidirectional_rnn case
+// (without cross links).
+constexpr int kAuxInputTensor = 9;       // Optional.
+constexpr int kFwAuxWeightsTensor = 10;  // Optional.
+constexpr int kBwAuxWeightsTensor = 11;  // Optional.
+// Output tensors.
+constexpr int kFwOutputTensor = 0;
+constexpr int kBwOutputTensor = 1;  // Only if merge_outputs is false.
+
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantize_weights.cc)
+
+// Temporary tensors.
+enum TemporaryTensor {
+  kInputQuantized = 0,
+  kFwHiddenStateQuantized = 1,
+  kBwHiddenStateQuantized = 2,
+  kScalingFactors = 3,
+  kAccumScratch = 4,
+  kZeroPoints = 5,
+  kFwRowSums = 6,
+  kBwRowSums = 7,
+  kAuxInputQuantized = 8,
+  kNumTemporaryTensors = 9
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  context->AddTensors(context, kNumTemporaryTensors,
+                      &op_data->scratch_tensor_index);
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  const auto* params = reinterpret_cast<TfLiteBidirectionalSequenceRNNParams*>(
+      node->builtin_data);
+
+  // Check we have all the inputs and outputs we need.
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 12);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size,
+                    params->merge_outputs ? 1 : 2);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* fw_input_weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kFwWeightsTensor,
+                                          &fw_input_weights));
+  const TfLiteTensor* fw_recurrent_weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFwRecurrentWeightsTensor,
+                                 &fw_recurrent_weights));
+  const TfLiteTensor* fw_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFwBiasTensor, &fw_bias));
+  const TfLiteTensor* fw_hidden_state;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kFwHiddenStateTensor,
+                                          &fw_hidden_state));
+  const TfLiteTensor* bw_input_weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBwWeightsTensor,
+                                          &bw_input_weights));
+  const TfLiteTensor* bw_recurrent_weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kBwRecurrentWeightsTensor,
+                                 &bw_recurrent_weights));
+  const TfLiteTensor* bw_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kBwBiasTensor, &bw_bias));
+  const TfLiteTensor* bw_hidden_state;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBwHiddenStateTensor,
+                                          &bw_hidden_state));
+
+  const TfLiteTensor* aux_input =
+      GetOptionalInputTensor(context, node, kAuxInputTensor);
+  const TfLiteTensor* fw_aux_input_weights =
+      GetOptionalInputTensor(context, node, kFwAuxWeightsTensor);
+  const TfLiteTensor* bw_aux_input_weights =
+      GetOptionalInputTensor(context, node, kBwAuxWeightsTensor);
+
+  const bool aux_inputs_weights_or_none =
+      ((fw_aux_input_weights != nullptr) &&
+       (bw_aux_input_weights != nullptr)) ||
+      ((fw_aux_input_weights == nullptr) && (bw_aux_input_weights == nullptr));
+  TF_LITE_ENSURE(context, aux_inputs_weights_or_none);
+  const bool has_aux_input = (fw_aux_input_weights != nullptr);
+
+  // Check all the parameters of tensor match within themselves and match the
+  // input configuration.
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+
+  TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
+  const bool time_major = params->time_major;
+  const int batch_size =
+      (time_major) ? input->dims->data[1] : input->dims->data[0];
+  const int max_time =
+      (time_major) ? input->dims->data[0] : input->dims->data[1];
+  const int fw_num_units = fw_input_weights->dims->data[0];
+  const int bw_num_units = bw_input_weights->dims->data[0];
+  TF_LITE_ENSURE_EQ(context, input->dims->data[2],
+                    fw_input_weights->dims->data[1]);
+  TF_LITE_ENSURE_EQ(context, input->dims->data[2],
+                    bw_input_weights->dims->data[1]);
+  TF_LITE_ENSURE_EQ(context, fw_input_weights->dims->data[0],
+                    fw_bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, bw_input_weights->dims->data[0],
+                    bw_bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, fw_recurrent_weights->dims->data[0],
+                    fw_bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, bw_recurrent_weights->dims->data[1],
+                    bw_bias->dims->data[0]);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(fw_hidden_state), 2);
+  TF_LITE_ENSURE_EQ(context, fw_hidden_state->dims->data[0], batch_size);
+  TF_LITE_ENSURE_EQ(context, fw_hidden_state->dims->data[1], fw_num_units);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(bw_hidden_state), 2);
+  TF_LITE_ENSURE_EQ(context, bw_hidden_state->dims->data[0], batch_size);
+  TF_LITE_ENSURE_EQ(context, bw_hidden_state->dims->data[1], bw_num_units);
+
+  if (has_aux_input) {
+    // Check that aux_input has the same dimensions (except last) as the input.
+    TF_LITE_ASSERT_EQ(aux_input->dims->data[0], input->dims->data[0]);
+    TF_LITE_ASSERT_EQ(aux_input->dims->data[1], input->dims->data[1]);
+    // Check that aux_input_weights has the same dimensions (except last) as
+    // the input_weights.
+    TF_LITE_ASSERT_EQ(fw_aux_input_weights->dims->data[0], fw_num_units);
+    TF_LITE_ASSERT_EQ(bw_aux_input_weights->dims->data[0], bw_num_units);
+    TF_LITE_ASSERT_EQ(aux_input->dims->data[2],
+                      fw_aux_input_weights->dims->data[1]);
+    TF_LITE_ASSERT_EQ(aux_input->dims->data[2],
+                      bw_aux_input_weights->dims->data[1]);
+  }
+
+  if (IsHybridOp(input, fw_input_weights)) {
+    OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+    op_data->fw_compute_row_sums = true;
+    op_data->bw_compute_row_sums = true;
+    TfLiteIntArrayFree(node->temporaries);
+    if (has_aux_input) {
+      node->temporaries = TfLiteIntArrayCreate(kNumTemporaryTensors);
+    } else {
+      // No need to create a temporary tensor for the non-existent aux_input.
+      node->temporaries = TfLiteIntArrayCreate(kNumTemporaryTensors - 1);
+    }
+
+    node->temporaries->data[kInputQuantized] =
+        op_data->scratch_tensor_index + kInputQuantized;
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kInputQuantized,
+                                                &input_quantized));
+    input_quantized->type = fw_input_weights->type;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
+      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                       input_quantized_size));
+    }
+
+    node->temporaries->data[kFwHiddenStateQuantized] =
+        op_data->scratch_tensor_index + kFwHiddenStateQuantized;
+    TfLiteTensor* fw_hidden_state_quantized;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, kFwHiddenStateQuantized,
+                                       &fw_hidden_state_quantized));
+    fw_hidden_state_quantized->type = fw_input_weights->type;
+    fw_hidden_state_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(fw_hidden_state_quantized->dims,
+                             fw_hidden_state->dims)) {
+      TfLiteIntArray* fw_hidden_state_quantized_size =
+          TfLiteIntArrayCopy(fw_hidden_state->dims);
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, fw_hidden_state_quantized,
+                                         fw_hidden_state_quantized_size));
+    }
+
+    node->temporaries->data[kBwHiddenStateQuantized] =
+        op_data->scratch_tensor_index + kBwHiddenStateQuantized;
+    TfLiteTensor* bw_hidden_state_quantized;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, kBwHiddenStateQuantized,
+                                       &bw_hidden_state_quantized));
+    bw_hidden_state_quantized->type = fw_input_weights->type;
+    bw_hidden_state_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(bw_hidden_state_quantized->dims,
+                             bw_hidden_state->dims)) {
+      TfLiteIntArray* bw_hidden_state_quantized_size =
+          TfLiteIntArrayCopy(bw_hidden_state->dims);
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, bw_hidden_state_quantized,
+                                         bw_hidden_state_quantized_size));
+    }
+
+    // Allocate temporary tensors to store scaling factors of quantization.
+    node->temporaries->data[kScalingFactors] =
+        op_data->scratch_tensor_index + kScalingFactors;
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kScalingFactors,
+                                                &scaling_factors));
+    scaling_factors->type = kTfLiteFloat32;
+    scaling_factors->allocation_type = kTfLiteArenaRw;
+    int scaling_dims[1] = {batch_size};
+    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
+      TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1);
+      scaling_factors_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors,
+                                                       scaling_factors_size));
+    }
+    node->temporaries->data[kAccumScratch] =
+        op_data->scratch_tensor_index + kAccumScratch;
+    TfLiteTensor* accum_scratch;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kAccumScratch,
+                                                &accum_scratch));
+    accum_scratch->type = kTfLiteInt32;
+    accum_scratch->allocation_type = kTfLiteArenaRw;
+    int accum_scratch_dims[2] = {std::max(fw_num_units, bw_num_units),
+                                 batch_size};
+    if (!TfLiteIntArrayEqualsArray(accum_scratch->dims, 2,
+                                   accum_scratch_dims)) {
+      TfLiteIntArray* accum_scratch_size = TfLiteIntArrayCreate(2);
+      accum_scratch_size->data[0] = accum_scratch_dims[0];
+      accum_scratch_size->data[1] = accum_scratch_dims[1];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, accum_scratch,
+                                                       accum_scratch_size));
+    }
+    node->temporaries->data[kZeroPoints] =
+        op_data->scratch_tensor_index + kZeroPoints;
+    TfLiteTensor* zero_points;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node, /*index=*/kZeroPoints, &zero_points));
+    zero_points->type = kTfLiteInt32;
+    zero_points->allocation_type = kTfLiteArenaRw;
+    int zero_points_dims[1] = {batch_size};
+    if (!TfLiteIntArrayEqualsArray(zero_points->dims, 1, zero_points_dims)) {
+      TfLiteIntArray* zero_points_size = TfLiteIntArrayCreate(1);
+      zero_points_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, zero_points,
+                                                       zero_points_size));
+    }
+    const int num_row_sums = has_aux_input ? 3 : 2;
+    node->temporaries->data[kFwRowSums] =
+        op_data->scratch_tensor_index + kFwRowSums;
+    TfLiteTensor* fw_row_sums;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node, /*index=*/kFwRowSums, &fw_row_sums));
+    fw_row_sums->type = kTfLiteInt32;
+    fw_row_sums->name = "Lstm_fw_row_sums";
+    fw_row_sums->allocation_type = kTfLiteArenaRwPersistent;
+    int fw_row_sums_dims[2] = {num_row_sums, fw_num_units};
+    if (!TfLiteIntArrayEqualsArray(fw_row_sums->dims, 2, fw_row_sums_dims)) {
+      TfLiteIntArray* fw_row_sums_size = TfLiteIntArrayCreate(2);
+      fw_row_sums_size->data[0] = fw_row_sums_dims[0];
+      fw_row_sums_size->data[1] = fw_row_sums_dims[1];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, fw_row_sums,
+                                                       fw_row_sums_size));
+    }
+    node->temporaries->data[kBwRowSums] =
+        op_data->scratch_tensor_index + kBwRowSums;
+    TfLiteTensor* bw_row_sums;
+    TF_LITE_ENSURE_OK(
+        context,
+        GetTemporarySafe(context, node, /*index=*/kBwRowSums, &bw_row_sums));
+    bw_row_sums->type = kTfLiteInt32;
+    bw_row_sums->name = "Lstm_bw_row_sums";
+    bw_row_sums->allocation_type = kTfLiteArenaRwPersistent;
+    int bw_row_sums_dims[2] = {num_row_sums, bw_num_units};
+    if (!TfLiteIntArrayEqualsArray(bw_row_sums->dims, 2, bw_row_sums_dims)) {
+      TfLiteIntArray* bw_row_sums_size = TfLiteIntArrayCreate(2);
+      bw_row_sums_size->data[0] = bw_row_sums_dims[0];
+      bw_row_sums_size->data[1] = bw_row_sums_dims[1];
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_row_sums,
+                                                       bw_row_sums_size));
+    }
+    if (has_aux_input) {
+      node->temporaries->data[kAuxInputQuantized] =
+          op_data->scratch_tensor_index + kAuxInputQuantized;
+      TfLiteTensor* aux_input_quantized;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, kAuxInputQuantized,
+                                         &aux_input_quantized));
+      aux_input_quantized->type = fw_input_weights->type;
+      aux_input_quantized->allocation_type = kTfLiteArenaRw;
+      if (!TfLiteIntArrayEqual(aux_input_quantized->dims, aux_input->dims)) {
+        TfLiteIntArray* aux_input_quantized_size =
+            TfLiteIntArrayCopy(aux_input->dims);
+        TF_LITE_ENSURE_OK(context,
+                          context->ResizeTensor(context, aux_input_quantized,
+                                                aux_input_quantized_size));
+      }
+    }
+  }
+
+  // Resize outputs.
+  TfLiteTensor* fw_output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFwOutputTensor, &fw_output));
+  TfLiteIntArray* fw_output_size_array = TfLiteIntArrayCreate(3);
+  fw_output_size_array->data[0] = (time_major) ? max_time : batch_size;
+  fw_output_size_array->data[1] = (time_major) ? batch_size : max_time;
+  fw_output_size_array->data[2] =
+      params->merge_outputs ? fw_num_units + bw_num_units : fw_num_units;
+  TF_LITE_ENSURE_OK(
+      context, context->ResizeTensor(context, fw_output, fw_output_size_array));
+  if (!params->merge_outputs) {
+    TfLiteTensor* bw_output;
+    TF_LITE_ENSURE_OK(
+        context, GetOutputSafe(context, node, kBwOutputTensor, &bw_output));
+    TfLiteIntArray* bw_output_size_array = TfLiteIntArrayCreate(3);
+    bw_output_size_array->data[0] = (time_major) ? max_time : batch_size;
+    bw_output_size_array->data[1] = (time_major) ? batch_size : max_time;
+    bw_output_size_array->data[2] = bw_num_units;
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, bw_output,
+                                                     bw_output_size_array));
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalFloat(const TfLiteTensor* input, const TfLiteTensor* bw_input,
+                       const TfLiteTensor* fw_input_weights,
+                       const TfLiteTensor* fw_recurrent_weights,
+                       const TfLiteTensor* fw_bias,
+                       const TfLiteTensor* bw_input_weights,
+                       const TfLiteTensor* bw_recurrent_weights,
+                       const TfLiteTensor* bw_bias,
+                       const TfLiteTensor* aux_input,
+                       const TfLiteTensor* fw_aux_input_weights,
+                       const TfLiteTensor* bw_aux_input_weights,
+                       const TfLiteBidirectionalSequenceRNNParams* params,
+                       TfLiteTensor* fw_hidden_state, TfLiteTensor* fw_output,
+                       TfLiteTensor* bw_hidden_state, TfLiteTensor* bw_output) {
+  const bool time_major = params->time_major;
+  const int batch_size =
+      (time_major) ? input->dims->data[1] : input->dims->data[0];
+  const int max_time =
+      (time_major) ? input->dims->data[0] : input->dims->data[1];
+  const int input_size = input->dims->data[2];
+  const int aux_input_size = (aux_input) ? aux_input->dims->data[2] : 0;
+
+  const int fw_num_units = fw_input_weights->dims->data[0];
+  const float* fw_bias_ptr = GetTensorData<float>(fw_bias);
+  const float* fw_input_weights_ptr = GetTensorData<float>(fw_input_weights);
+  const float* fw_recurrent_weights_ptr =
+      GetTensorData<float>(fw_recurrent_weights);
+
+  const int bw_num_units = bw_input_weights->dims->data[0];
+  const float* bw_bias_ptr = GetTensorData<float>(bw_bias);
+  const float* bw_input_weights_ptr = GetTensorData<float>(bw_input_weights);
+  const float* bw_recurrent_weights_ptr =
+      GetTensorData<float>(bw_recurrent_weights);
+
+  const float* fw_aux_input_weights_ptr =
+      (fw_aux_input_weights != nullptr)
+          ? GetTensorData<float>(fw_aux_input_weights)
+          : nullptr;
+  const float* bw_aux_input_weights_ptr =
+      (bw_aux_input_weights != nullptr)
+          ? GetTensorData<float>(bw_aux_input_weights)
+          : nullptr;
+
+  const int fw_output_step =
+      params->merge_outputs ? fw_num_units + bw_num_units : fw_num_units;
+  const int bw_output_step =
+      params->merge_outputs ? fw_num_units + bw_num_units : bw_num_units;
+  if (time_major) {
+    // Forward cell.
+    float* fw_hidden_state_ptr_batch = GetTensorData<float>(fw_hidden_state);
+    for (int s = 0; s < max_time; s++) {
+      const float* input_ptr_batch =
+          GetTensorData<float>(input) + s * input_size * batch_size;
+      const float* aux_input_ptr_batch =
+          (aux_input != nullptr)
+              ? GetTensorData<float>(aux_input) + s * input_size * batch_size
+              : nullptr;
+      float* output_ptr_batch =
+          GetTensorData<float>(fw_output) + s * fw_output_step * batch_size;
+
+      kernel_utils::RnnBatchStep(
+          input_ptr_batch, fw_input_weights_ptr, aux_input_ptr_batch,
+          fw_aux_input_weights_ptr, fw_recurrent_weights_ptr, fw_bias_ptr,
+          input_size, aux_input_size, fw_num_units, batch_size, fw_output_step,
+          params->activation, fw_hidden_state_ptr_batch, output_ptr_batch);
+    }
+    // Backward cell.
+    float* bw_hidden_state_ptr_batch = GetTensorData<float>(bw_hidden_state);
+    for (int s = max_time - 1; s >= 0; s--) {
+      const float* input_ptr_batch =
+          GetTensorData<float>(bw_input) + s * input_size * batch_size;
+      const float* aux_input_ptr_batch =
+          (aux_input != nullptr)
+              ? GetTensorData<float>(aux_input) + s * input_size * batch_size
+              : nullptr;
+      float* output_ptr_batch =
+          (params->merge_outputs
+               ? GetTensorData<float>(fw_output) + fw_num_units
+               : GetTensorData<float>(bw_output)) +
+          s * bw_output_step * batch_size;
+
+      kernel_utils::RnnBatchStep(
+          input_ptr_batch, bw_input_weights_ptr, aux_input_ptr_batch,
+          bw_aux_input_weights_ptr, bw_recurrent_weights_ptr, bw_bias_ptr,
+          input_size, aux_input_size, bw_num_units, batch_size, bw_output_step,
+          params->activation, bw_hidden_state_ptr_batch, output_ptr_batch);
+    }
+  } else {
+    for (int b = 0; b < batch_size; b++) {
+      // Forward cell.
+      float* fw_hidden_state_ptr_batch =
+          GetTensorData<float>(fw_hidden_state) + b * fw_num_units;
+      float* fw_output_offset =
+          GetTensorData<float>(fw_output) + b * fw_output_step * max_time;
+      for (int s = 0; s < max_time; s++) {
+        const float* input_ptr_batch = GetTensorData<float>(input) +
+                                       b * input_size * max_time +
+                                       s * input_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) +
+                      b * aux_input_size * max_time + s * aux_input_size
+                : nullptr;
+        float* output_ptr_batch = fw_output_offset + s * fw_output_step;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, fw_input_weights_ptr, aux_input_ptr_batch,
+            fw_aux_input_weights_ptr, fw_recurrent_weights_ptr, fw_bias_ptr,
+            input_size, aux_input_size, fw_num_units, /*batch_size=*/1,
+            fw_output_step, params->activation, fw_hidden_state_ptr_batch,
+            output_ptr_batch);
+      }
+      // Backward cell.
+      float* bw_hidden_state_ptr_batch =
+          GetTensorData<float>(bw_hidden_state) + b * bw_num_units;
+      float* bw_output_offset =
+          params->merge_outputs
+              ? GetTensorData<float>(fw_output) +
+                    b * bw_output_step * max_time + fw_num_units
+              : GetTensorData<float>(bw_output) + b * bw_output_step * max_time;
+      for (int s = max_time - 1; s >= 0; s--) {
+        const float* input_ptr_batch = GetTensorData<float>(input) +
+                                       b * input_size * max_time +
+                                       s * input_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) +
+                      b * aux_input_size * max_time + s * aux_input_size
+                : nullptr;
+        float* output_ptr_batch = bw_output_offset + s * bw_output_step;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, bw_input_weights_ptr, aux_input_ptr_batch,
+            bw_aux_input_weights_ptr, bw_recurrent_weights_ptr, bw_bias_ptr,
+            input_size, aux_input_size, bw_num_units, /*batch_size=*/1,
+            bw_output_step, params->activation, bw_hidden_state_ptr_batch,
+            output_ptr_batch);
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalHybrid(
+    const TfLiteTensor* input, const TfLiteTensor* bw_input,
+    const TfLiteTensor* fw_input_weights,
+    const TfLiteTensor* fw_recurrent_weights, const TfLiteTensor* fw_bias,
+    const TfLiteTensor* bw_input_weights,
+    const TfLiteTensor* bw_recurrent_weights, const TfLiteTensor* bw_bias,
+    const TfLiteTensor* aux_input, const TfLiteTensor* aux_fw_input_weights,
+    const TfLiteTensor* aux_bw_input_weights,
+    const TfLiteBidirectionalSequenceRNNParams* params,
+    TfLiteTensor* scaling_factors, TfLiteTensor* input_quantized,
+    TfLiteTensor* aux_input_quantized, TfLiteTensor* fw_hidden_state_quantized,
+    TfLiteTensor* fw_hidden_state, TfLiteTensor* fw_output,
+    TfLiteTensor* bw_hidden_state_quantized, TfLiteTensor* bw_hidden_state,
+    TfLiteTensor* bw_output, TfLiteTensor* zero_points,
+    TfLiteTensor* accum_scratch, TfLiteTensor* fw_row_sums,
+    TfLiteTensor* bw_row_sums, bool* fw_compute_row_sums,
+    bool* bw_compute_row_sums) {
+  const bool time_major = params->time_major;
+  const int batch_size =
+      (time_major) ? input->dims->data[1] : input->dims->data[0];
+  const int max_time =
+      (time_major) ? input->dims->data[0] : input->dims->data[1];
+  const int input_size = input->dims->data[2];
+  const int aux_input_size = (aux_input) ? aux_input->dims->data[2] : 0;
+
+  const int fw_num_units = fw_input_weights->dims->data[0];
+  const float* fw_bias_ptr = GetTensorData<float>(fw_bias);
+  const int8_t* fw_input_weights_ptr = GetTensorData<int8_t>(fw_input_weights);
+  float fw_input_weights_scale = fw_input_weights->params.scale;
+  const int8_t* fw_recurrent_weights_ptr =
+      GetTensorData<int8_t>(fw_recurrent_weights);
+  float fw_recurrent_weights_scale = fw_recurrent_weights->params.scale;
+
+  const int bw_num_units = bw_input_weights->dims->data[0];
+  const float* bw_bias_ptr = GetTensorData<float>(bw_bias);
+  const int8_t* bw_input_weights_ptr = GetTensorData<int8_t>(bw_input_weights);
+  float bw_input_weights_scale = bw_input_weights->params.scale;
+  const int8_t* bw_recurrent_weights_ptr =
+      GetTensorData<int8_t>(bw_recurrent_weights);
+  float bw_recurrent_weights_scale = bw_recurrent_weights->params.scale;
+
+  // Set the auxiliary pointers and scales if needed.
+  const int8_t* aux_fw_input_weights_ptr = nullptr;
+  float aux_fw_input_weights_scale = 0.0f;
+  const int8_t* aux_bw_input_weights_ptr = nullptr;
+  float aux_bw_input_weights_scale = 0.0f;
+  int8_t* aux_quantized_input_ptr = nullptr;
+  if (aux_input_size > 0) {
+    aux_fw_input_weights_ptr = GetTensorData<int8_t>(aux_fw_input_weights);
+    aux_fw_input_weights_scale = aux_fw_input_weights->params.scale;
+    aux_bw_input_weights_ptr = GetTensorData<int8_t>(aux_bw_input_weights);
+    aux_bw_input_weights_scale = aux_bw_input_weights->params.scale;
+    aux_quantized_input_ptr = GetTensorData<int8_t>(aux_input_quantized);
+  }
+
+  // Initialize temporary storage for quantized values.
+  int8_t* quantized_input_ptr = GetTensorData<int8_t>(input_quantized);
+  int8_t* fw_quantized_hidden_state_ptr =
+      GetTensorData<int8_t>(fw_hidden_state_quantized);
+  int8_t* bw_quantized_hidden_state_ptr =
+      GetTensorData<int8_t>(bw_hidden_state_quantized);
+  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors);
+  int32_t* accum_scratch_ptr = GetTensorData<int32_t>(accum_scratch);
+  int32_t* zero_points_ptr = nullptr;
+  int32_t* fw_row_sums_ptr = nullptr;
+  int32_t* bw_row_sums_ptr = nullptr;
+  if (params->asymmetric_quantize_inputs) {
+    zero_points_ptr = GetTensorData<int32_t>(zero_points);
+    fw_row_sums_ptr = GetTensorData<int32_t>(fw_row_sums);
+    bw_row_sums_ptr = GetTensorData<int32_t>(bw_row_sums);
+  }
+  const int fw_output_step =
+      params->merge_outputs ? fw_num_units + bw_num_units : fw_num_units;
+  const int bw_output_step =
+      params->merge_outputs ? fw_num_units + bw_num_units : bw_num_units;
+
+  if (time_major) {
+    for (int t = 0; t < max_time; t++) {
+      // Forward cell.
+      float* fw_hidden_state_ptr_batch = GetTensorData<float>(fw_hidden_state);
+      for (int s = 0; s < max_time; s++) {
+        const float* input_ptr_batch =
+            GetTensorData<float>(input) + s * input_size * batch_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) + s * input_size * batch_size
+                : nullptr;
+        float* output_ptr_batch =
+            GetTensorData<float>(fw_output) + s * fw_output_step * batch_size;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, fw_input_weights_ptr, fw_input_weights_scale,
+            aux_input_ptr_batch, aux_fw_input_weights_ptr,
+            aux_fw_input_weights_scale, fw_recurrent_weights_ptr,
+            fw_recurrent_weights_scale, fw_bias_ptr, input_size, aux_input_size,
+            fw_num_units, batch_size, fw_output_step, params->activation,
+            quantized_input_ptr, aux_quantized_input_ptr,
+            fw_quantized_hidden_state_ptr, scaling_factors_ptr,
+            fw_hidden_state_ptr_batch, output_ptr_batch,
+            params->asymmetric_quantize_inputs, zero_points_ptr,
+            accum_scratch_ptr, fw_row_sums_ptr, fw_compute_row_sums);
+      }
+      // Backward cell.
+      float* bw_hidden_state_ptr_batch = GetTensorData<float>(bw_hidden_state);
+      for (int s = max_time - 1; s >= 0; s--) {
+        const float* input_ptr_batch =
+            GetTensorData<float>(bw_input) + s * input_size * batch_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) + s * input_size * batch_size
+                : nullptr;
+        float* output_ptr_batch =
+            (params->merge_outputs
+                 ? GetTensorData<float>(fw_output) + fw_num_units
+                 : GetTensorData<float>(bw_output)) +
+            s * bw_output_step * batch_size;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, bw_input_weights_ptr, bw_input_weights_scale,
+            aux_input_ptr_batch, aux_bw_input_weights_ptr,
+            aux_bw_input_weights_scale, bw_recurrent_weights_ptr,
+            bw_recurrent_weights_scale, bw_bias_ptr, input_size, aux_input_size,
+            bw_num_units, batch_size, bw_output_step, params->activation,
+            quantized_input_ptr, aux_quantized_input_ptr,
+            bw_quantized_hidden_state_ptr, scaling_factors_ptr,
+            bw_hidden_state_ptr_batch, output_ptr_batch,
+            params->asymmetric_quantize_inputs, zero_points_ptr,
+            accum_scratch_ptr, bw_row_sums_ptr, bw_compute_row_sums);
+      }
+    }
+  } else {
+    for (int b = 0; b < batch_size; b++) {
+      // Forward cell.
+      float* fw_hidden_state_ptr_batch =
+          GetTensorData<float>(fw_hidden_state) + b * fw_num_units;
+      float* fw_output_offset =
+          GetTensorData<float>(fw_output) + b * fw_output_step * max_time;
+      for (int s = 0; s < max_time; s++) {
+        const float* input_ptr_batch = GetTensorData<float>(input) +
+                                       b * input_size * max_time +
+                                       s * input_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) + b * input_size * max_time +
+                      s * input_size
+                : nullptr;
+        float* output_ptr_batch = fw_output_offset + s * fw_output_step;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, fw_input_weights_ptr, fw_input_weights_scale,
+            aux_input_ptr_batch, aux_fw_input_weights_ptr,
+            aux_fw_input_weights_scale, fw_recurrent_weights_ptr,
+            fw_recurrent_weights_scale, fw_bias_ptr, input_size, aux_input_size,
+            fw_num_units, /*batch_size=*/1, fw_output_step, params->activation,
+            quantized_input_ptr, aux_quantized_input_ptr,
+            fw_quantized_hidden_state_ptr, scaling_factors_ptr,
+            fw_hidden_state_ptr_batch, output_ptr_batch,
+            params->asymmetric_quantize_inputs, zero_points_ptr,
+            accum_scratch_ptr, fw_row_sums_ptr, fw_compute_row_sums);
+      }
+      // Backward cell.
+      float* bw_hidden_state_ptr_batch =
+          GetTensorData<float>(bw_hidden_state) + b * bw_num_units;
+      float* bw_output_offset =
+          params->merge_outputs
+              ? GetTensorData<float>(fw_output) +
+                    b * bw_output_step * max_time + fw_num_units
+              : GetTensorData<float>(bw_output) + b * bw_output_step * max_time;
+      for (int s = max_time - 1; s >= 0; s--) {
+        const float* input_ptr_batch = GetTensorData<float>(input) +
+                                       b * input_size * max_time +
+                                       s * input_size;
+        const float* aux_input_ptr_batch =
+            (aux_input != nullptr)
+                ? GetTensorData<float>(aux_input) + b * input_size * max_time +
+                      s * input_size
+                : nullptr;
+        float* output_ptr_batch = bw_output_offset + s * bw_output_step;
+
+        kernel_utils::RnnBatchStep(
+            input_ptr_batch, bw_input_weights_ptr, bw_input_weights_scale,
+            aux_input_ptr_batch, aux_bw_input_weights_ptr,
+            aux_bw_input_weights_scale, bw_recurrent_weights_ptr,
+            bw_recurrent_weights_scale, bw_bias_ptr, input_size, aux_input_size,
+            bw_num_units, /*batch_size=*/1, bw_output_step, params->activation,
+            quantized_input_ptr, aux_quantized_input_ptr,
+            bw_quantized_hidden_state_ptr, scaling_factors_ptr,
+            bw_hidden_state_ptr_batch, output_ptr_batch,
+            params->asymmetric_quantize_inputs, zero_points_ptr,
+            accum_scratch_ptr, bw_row_sums_ptr, bw_compute_row_sums);
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const auto* params = reinterpret_cast<TfLiteBidirectionalSequenceRNNParams*>(
+      node->builtin_data);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* fw_input_weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kFwWeightsTensor,
+                                          &fw_input_weights));
+  const TfLiteTensor* fw_recurrent_weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFwRecurrentWeightsTensor,
+                                 &fw_recurrent_weights));
+  const TfLiteTensor* fw_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFwBiasTensor, &fw_bias));
+  const TfLiteTensor* bw_input_weights;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBwWeightsTensor,
+                                          &bw_input_weights));
+  const TfLiteTensor* bw_recurrent_weights;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kBwRecurrentWeightsTensor,
+                                 &bw_recurrent_weights));
+  const TfLiteTensor* bw_bias;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kBwBiasTensor, &bw_bias));
+
+  // Get auxiliary inputs.
+  const TfLiteTensor* aux_input =
+      GetOptionalInputTensor(context, node, kAuxInputTensor);
+  const TfLiteTensor* fw_aux_input_weights =
+      GetOptionalInputTensor(context, node, kFwAuxWeightsTensor);
+  const TfLiteTensor* bw_aux_input_weights =
+      GetOptionalInputTensor(context, node, kBwAuxWeightsTensor);
+
+  TfLiteTensor* fw_hidden_state =
+      GetVariableInput(context, node, kFwHiddenStateTensor);
+  TFLITE_DCHECK(fw_hidden_state != nullptr);
+  TfLiteTensor* bw_hidden_state =
+      GetVariableInput(context, node, kBwHiddenStateTensor);
+  TFLITE_DCHECK(bw_hidden_state != nullptr);
+
+  TfLiteTensor* fw_output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kFwOutputTensor, &fw_output));
+  TfLiteTensor* bw_output = params->merge_outputs
+                                ? nullptr
+                                : GetOutput(context, node, kBwOutputTensor);
+
+  const bool has_previous_bw_output = (aux_input != nullptr);
+  const bool use_aux_input = (fw_aux_input_weights != nullptr);
+
+  // We want to cover the following cases:
+  //
+  // If not stacking (not connected after other bidi lstms):
+  //   both fw & bw will just use `input`; aux_input will be null.
+  //
+  // If stacking with cross_links, TensorFlow equivalent
+  // (tf.contrib.rnn.stack_bidirectional_rnn):
+  //   both fw & bw will use `input`, but aux_input will be none null.
+  //   Note, this time, whether connected after other bidi lstms both works.
+  //
+  // If stacking without cross_links, but connected after other bidi lstms,
+  // TensorFlow equivalent (tf.nn.static_bidirectional_rnn):
+  //   fw will use `input`, bw will use aux_input, and the `real aux_input`
+  //   will be null.
+
+  const bool non_stacking_mode = !use_aux_input && has_previous_bw_output;
+  const TfLiteTensor* bw_input = non_stacking_mode ? aux_input : input;
+  const TfLiteTensor* real_aux_input = non_stacking_mode ? nullptr : aux_input;
+
+  switch (fw_input_weights->type) {
+    case kTfLiteFloat32:
+      return EvalFloat(input, bw_input, fw_input_weights, fw_recurrent_weights,
+                       fw_bias, bw_input_weights, bw_recurrent_weights, bw_bias,
+                       real_aux_input, fw_aux_input_weights,
+                       bw_aux_input_weights, params, fw_hidden_state, fw_output,
+                       bw_hidden_state, bw_output);
+    case kTfLiteUInt8:
+    case kTfLiteInt8: {
+      TfLiteTensor* input_quantized;
+      TF_LITE_ENSURE_OK(
+          context,
+          GetTemporarySafe(context, node, kInputQuantized, &input_quantized));
+      TfLiteTensor* fw_hidden_state_quantized;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, kFwHiddenStateQuantized,
+                                         &fw_hidden_state_quantized));
+      TfLiteTensor* bw_hidden_state_quantized;
+      TF_LITE_ENSURE_OK(context,
+                        GetTemporarySafe(context, node, kBwHiddenStateQuantized,
+                                         &bw_hidden_state_quantized));
+      TfLiteTensor* scaling_factors;
+      TF_LITE_ENSURE_OK(
+          context,
+          GetTemporarySafe(context, node, kScalingFactors, &scaling_factors));
+      TfLiteTensor* zero_points;
+      TF_LITE_ENSURE_OK(
+          context, GetTemporarySafe(context, node, kZeroPoints, &zero_points));
+      TfLiteTensor* accum_scratch;
+      TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kAccumScratch,
+                                                  &accum_scratch));
+      TfLiteTensor* fw_row_sums;
+      TF_LITE_ENSURE_OK(
+          context, GetTemporarySafe(context, node, kFwRowSums, &fw_row_sums));
+      TfLiteTensor* bw_row_sums;
+      TF_LITE_ENSURE_OK(
+          context, GetTemporarySafe(context, node, kBwRowSums, &bw_row_sums));
+      TfLiteTensor* aux_input_quantized =
+          use_aux_input ? GetTemporary(context, node, kAuxInputQuantized)
+                        : nullptr;
+      auto* op_data = reinterpret_cast<OpData*>(node->user_data);
+      return EvalHybrid(
+          input, bw_input, fw_input_weights, fw_recurrent_weights, fw_bias,
+          bw_input_weights, bw_recurrent_weights, bw_bias, real_aux_input,
+          fw_aux_input_weights, bw_aux_input_weights, params, scaling_factors,
+          input_quantized, aux_input_quantized, fw_hidden_state_quantized,
+          fw_hidden_state, fw_output, bw_hidden_state_quantized,
+          bw_hidden_state, bw_output, zero_points, accum_scratch, fw_row_sums,
+          bw_row_sums, &op_data->fw_compute_row_sums,
+          &op_data->bw_compute_row_sums);
+    }
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type not currently supported.");
+      return kTfLiteError;
+  }
+}
+
+}  // namespace bidirectional_sequence_rnn
+
+TfLiteRegistration* Register_BIDIRECTIONAL_SEQUENCE_RNN() {
+  static TfLiteRegistration r = {
+      bidirectional_sequence_rnn::Init, bidirectional_sequence_rnn::Free,
+      bidirectional_sequence_rnn::Prepare, bidirectional_sequence_rnn::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/bitcast.cc

@@ -0,0 +1,118 @@
+/* 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.
+==============================================================================*/
+
+#include <cstring>
+#include <memory>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace bitcast {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus CalculateShape(TfLiteContext* context, const TfLiteTensor* input,
+                            const TfLiteTensor* output,
+                            TfLiteIntArray** output_shape) {
+  const int dims = NumDimensions(input);
+
+  auto input_type = input->type;
+  auto output_type = output->type;
+  size_t input_type_size, output_type_size;
+  TF_LITE_ENSURE_STATUS(GetSizeOfType(context, input_type, &input_type_size));
+  TF_LITE_ENSURE_STATUS(GetSizeOfType(context, output_type, &output_type_size));
+
+  TfLiteIntArray* shape = nullptr;
+  if (input_type_size > output_type_size) {
+    // If the input datatype T is larger than the output datatype type then the
+    // shape changes from [...] to [..., sizeof(T)/sizeof(type)].
+    shape = TfLiteIntArrayCreate(dims + 1);
+    for (int i = 0; i < dims; ++i) {
+      shape->data[i] = input->dims->data[i];
+    }
+    shape->data[dims] = input_type_size / output_type_size;
+  } else if (input_type_size < output_type_size) {
+    // If T is smaller than type, the operator requires that the rightmost
+    // dimension be equal to sizeof(type)/sizeof(T). The shape then goes from
+    // [..., sizeof(type)/sizeof(T)] to [...].
+    TF_LITE_ENSURE_EQ(context, input->dims->data[dims - 1],
+                      output_type_size / input_type_size);
+    shape = TfLiteIntArrayCreate(dims - 1);
+    for (int i = 0; i < dims - 1; ++i) {
+      shape->data[i] = input->dims->data[i];
+    }
+  } else {
+    // Same element type size.
+    shape = TfLiteIntArrayCopy(input->dims);
+  }
+  *output_shape = shape;
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  TfLiteIntArray* output_size = nullptr;
+  TF_LITE_ENSURE_OK(context,
+                    CalculateShape(context, input, output, &output_size));
+  return context->ResizeTensor(context, output, output_size);
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // Only copy data if input and output do not share a buffer.
+  if (output->data.data != input->data.data) {
+    memcpy(output->data.data, input->data.data, input->bytes);
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace bitcast
+
+TfLiteRegistration* Register_BITCAST() {
+  static TfLiteRegistration r = {
+      nullptr,
+      nullptr,
+      bitcast::Prepare,
+      bitcast::Eval,
+      /*profiling_string=*/nullptr,
+      /*builtin_code=*/0,
+      /*custom_name=*/nullptr,
+      /*version=*/0,
+      /*registration_external=*/nullptr,
+      /*async_kernel=*/nullptr,
+      kTfLiteInplaceOpInput0Shared | kTfLiteInplaceOpDataUnmodified};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/bitcast_test.cc

@@ -0,0 +1,159 @@
+/* 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.
+==============================================================================*/
+
+#include <algorithm>
+#include <cstdint>
+#include <iterator>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+// std::bit_cast is only available since c++20, provide a handrolled version
+// of bit_cast here. Implementation modified from abseil bit_cast.
+template <
+    typename Dest, typename Source,
+    typename std::enable_if<sizeof(Dest) == sizeof(Source) &&
+                                std::is_trivially_copyable<Source>::value &&
+                                std::is_trivially_copyable<Dest>::value &&
+                                std::is_default_constructible<Dest>::value,
+                            int>::type = 0>
+inline Dest bit_cast(const Source& source) {
+  Dest dest;
+  memcpy(static_cast<void*>(std::addressof(dest)),
+         static_cast<const void*>(std::addressof(source)), sizeof(dest));
+  return dest;
+}
+
+class BitcastOpModel : public SingleOpModel {
+ public:
+  BitcastOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_BITCAST, BuiltinOptions_BitcastOptions,
+                 CreateBitcastOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input() const { return input_; }
+  int output() const { return output_; }
+
+ protected:
+  int input_;
+  int output_;
+};
+
+TEST(BitcastOpModel, BitcastInt32ToUint32) {
+  BitcastOpModel m({TensorType_INT32, {2, 3}}, {TensorType_UINT32, {2, 3}});
+  std::vector<int32_t> input = {INT32_MIN, -100, -1, 0, 100, INT32_MAX};
+  m.PopulateTensor<int32_t>(m.input(), input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<uint32_t> output;
+
+  std::transform(input.cbegin(), input.cend(), std::back_inserter(output),
+                 [](int32_t a) { return bit_cast<std::uint32_t>(a); });
+  EXPECT_THAT(m.ExtractVector<uint32_t>(m.output()), ElementsAreArray(output));
+}
+
+TEST(BitcastOpModel, BitcastUInt32ToInt32Inplace) {
+  BitcastOpModel m({TensorType_UINT32, {2, 3}}, {TensorType_INT32, {2, 3}});
+  std::vector<uint32_t> input = {0,
+                                 1,
+                                 100,
+                                 bit_cast<uint32_t>(INT32_MAX),
+                                 bit_cast<uint32_t>(INT32_MIN),
+                                 UINT32_MAX};
+  m.PopulateTensor<uint32_t>(m.input(), input);
+  const int kInplaceTensorIdx = 0;
+  const TfLiteTensor* input_tensor = m.GetInputTensor(kInplaceTensorIdx);
+  TfLiteTensor* output_tensor = m.GetOutputTensor(kInplaceTensorIdx);
+  output_tensor->data.data = input_tensor->data.data;
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<int32_t> output;
+
+  std::transform(input.cbegin(), input.cend(), std::back_inserter(output),
+                 [](uint32_t a) { return bit_cast<std::uint32_t>(a); });
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()), ElementsAreArray(output));
+  EXPECT_EQ(output_tensor->data.data, input_tensor->data.data);
+}
+
+TEST(BitcastOpModel, BitcastUInt32ToInt32) {
+  BitcastOpModel m({TensorType_UINT32, {2, 3}}, {TensorType_INT32, {2, 3}});
+  std::vector<uint32_t> input = {0,
+                                 1,
+                                 100,
+                                 bit_cast<uint32_t>(INT32_MAX),
+                                 bit_cast<uint32_t>(INT32_MIN),
+                                 UINT32_MAX};
+  m.PopulateTensor<uint32_t>(m.input(), input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<int32_t> output;
+
+  std::transform(input.cbegin(), input.cend(), std::back_inserter(output),
+                 [](uint32_t a) { return bit_cast<std::uint32_t>(a); });
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()), ElementsAreArray(output));
+}
+
+TEST(BitcastOpModel, BitcastUInt32Toint16) {
+  BitcastOpModel m({TensorType_UINT32, {2, 1}}, {TensorType_INT16, {2, 1, 2}});
+  std::vector<uint32_t> input = {(uint32_t)UINT16_MAX + 1,
+                                 (uint32_t)UINT16_MAX};
+  m.PopulateTensor<uint32_t>(m.input(), input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && \
+    __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+  // 00..01 00..00
+  // 00..00 11..11
+  std::vector<int16_t> output = {1, 0, 0, -1};
+#else
+  // 00..00 00..01
+  // 11..11 00..00
+  std::vector<int16_t> output = {0, 1, -1, 0};
+#endif
+  EXPECT_THAT(m.ExtractVector<int16_t>(m.output()), ElementsAreArray(output));
+}
+
+TEST(BitcastOpModel, BitcastInt16ToUint32) {
+  BitcastOpModel m({TensorType_INT16, {2, 1, 2}}, {TensorType_UINT32, {2, 1}});
+#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && \
+    __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+  std::vector<int16_t> input = {1, 0, 0, -1};
+#else
+  std::vector<int16_t> input = {0, 1, -1, 0};
+#endif
+  m.PopulateTensor<int16_t>(m.input(), input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<uint32_t> output = {(uint32_t)UINT16_MAX + 1,
+                                  (uint32_t)UINT16_MAX};
+  EXPECT_THAT(m.ExtractVector<uint32_t>(m.output()), ElementsAreArray(output));
+}
+
+TEST(BitcastOpModel, BitcastInt16ToUint32WrongShape) {
+#if GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(BitcastOpModel m({TensorType_INT16, {2, 2, 7}},
+                                {TensorType_UINT32, {2, 7}}),
+               "7 != 2");
+#endif
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/bitwise_xor.cc

@@ -0,0 +1,168 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace bitwise_xor {
+
+// Input/output tensor index.
+constexpr int kInputTensor1 = 0;
+constexpr int kInputTensor2 = 1;
+constexpr int kOutputTensor = 0;
+
+// Op data for bitwise xor op.
+struct OpData {
+  bool requires_broadcast = false;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* data = new OpData;
+  return data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
+
+  output->type = input1->type;
+
+  data->requires_broadcast = !HaveSameShapes(input1, input2);
+
+  TfLiteIntArray* output_size = nullptr;
+  if (data->requires_broadcast) {
+    TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
+                                   context, input1, input2, &output_size));
+  } else {
+    output_size = TfLiteIntArrayCopy(input1->dims);
+  }
+
+  return context->ResizeTensor(context, output, output_size);
+}
+
+template <typename T>
+T BitwiseXor(T x, T y) {
+  return x ^ y;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  const TfLiteType type = output->type;
+  switch (type) {
+    // The fallthrough is indended. Since bitwise xor function operates on the
+    // underlying binary representation of the integers, both integers and
+    // unsigned integers will have the same behavior
+    case kTfLiteUInt8:
+    case kTfLiteInt8: {
+      if (data->requires_broadcast) {
+        reference_ops::BroadcastBinaryFunction4DSlow<int8_t, int8_t, int8_t>(
+            GetTensorShape(input1), GetTensorData<int8_t>(input1),
+            GetTensorShape(input2), GetTensorData<int8_t>(input2),
+            GetTensorShape(output), GetTensorData<int8_t>(output), BitwiseXor);
+      } else {
+        reference_ops::BinaryFunction<int8_t, int8_t, int8_t>(
+            GetTensorShape(input1), GetTensorData<int8_t>(input1),
+            GetTensorShape(input2), GetTensorData<int8_t>(input2),
+            GetTensorShape(output), GetTensorData<int8_t>(output), BitwiseXor);
+      }
+      break;
+    }
+    case kTfLiteUInt16:
+    case kTfLiteInt16: {
+      if (data->requires_broadcast) {
+        reference_ops::BroadcastBinaryFunction4DSlow<int16_t, int16_t, int16_t>(
+            GetTensorShape(input1), GetTensorData<int16_t>(input1),
+            GetTensorShape(input2), GetTensorData<int16_t>(input2),
+            GetTensorShape(output), GetTensorData<int16_t>(output), BitwiseXor);
+      } else {
+        reference_ops::BinaryFunction<int16_t, int16_t, int16_t>(
+            GetTensorShape(input1), GetTensorData<int16_t>(input1),
+            GetTensorShape(input2), GetTensorData<int16_t>(input2),
+            GetTensorShape(output), GetTensorData<int16_t>(output), BitwiseXor);
+      }
+      break;
+    }
+    case kTfLiteUInt32:
+    case kTfLiteInt32: {
+      if (data->requires_broadcast) {
+        reference_ops::BroadcastBinaryFunction4DSlow<int32_t, int32_t, int32_t>(
+            GetTensorShape(input1), GetTensorData<int32_t>(input1),
+            GetTensorShape(input2), GetTensorData<int32_t>(input2),
+            GetTensorShape(output), GetTensorData<int32_t>(output), BitwiseXor);
+      } else {
+        reference_ops::BinaryFunction<int32_t, int32_t, int32_t>(
+            GetTensorShape(input1), GetTensorData<int32_t>(input1),
+            GetTensorShape(input2), GetTensorData<int32_t>(input2),
+            GetTensorShape(output), GetTensorData<int32_t>(output), BitwiseXor);
+      }
+      break;
+    }
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "BitwiseXor currently only supports "
+                         "8-bit/16-bit/32-bit integer/unsigned integer, got %s",
+                         TfLiteTypeGetName(type));
+      return kTfLiteError;
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace bitwise_xor
+
+TfLiteRegistration* Register_BITWISE_XOR() {
+  static TfLiteRegistration r = {bitwise_xor::Init, bitwise_xor::Free,
+                                 bitwise_xor::Prepare, bitwise_xor::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/bitwise_xor_test.cc

@@ -0,0 +1,140 @@
+/* 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.
+==============================================================================*/
+
+#include <cstdint>
+#include <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class BitwiseXorOpModel : public SingleOpModel {
+ public:
+  BitwiseXorOpModel(std::initializer_list<int> input1_shape,
+                    std::initializer_list<int> input2_shape,
+                    TensorType tensor_type) {
+    input1_ = AddInput(tensor_type);
+    input2_ = AddInput(tensor_type);
+    output_ = AddOutput(tensor_type);
+    SetBuiltinOp(BuiltinOperator_BITWISE_XOR, BuiltinOptions_BitwiseXorOptions,
+                 CreateBitwiseXorOptions(builder_).Union());
+    BuildInterpreter({input1_shape, input2_shape});
+  }
+
+  int input1() const { return input1_; }
+  int input2() const { return input2_; }
+
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  int input1_;
+  int input2_;
+  int output_;
+};
+
+TEST(BitwiseXorOpTest, SimpleTestInt8) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT8);
+  model.PopulateTensor<int8_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<int8_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<int8_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, SimpleTestInt16) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT16);
+  model.PopulateTensor<int16_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<int16_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<int16_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, SimpleTestInt32) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int32_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<int32_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<int32_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, SimpleTestUInt8) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_UINT8);
+  model.PopulateTensor<uint8_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<uint8_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<uint8_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, SimpleTestUInt16) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_UINT16);
+  model.PopulateTensor<uint16_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<uint16_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<uint16_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, SimpleTestUInt32) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_UINT32);
+  model.PopulateTensor<uint32_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<uint32_t>(model.input2(), {5, 0, 7, 11});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<uint32_t>(), ElementsAreArray({5, 5, 4, 5}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, BroadcastLhs) {
+  BitwiseXorOpModel model({1, 1, 1, 1}, {1, 1, 1, 4}, TensorType_INT32);
+  model.PopulateTensor<int32_t>(model.input1(), {5});
+  model.PopulateTensor<int32_t>(model.input2(), {0, -5, -3, 14});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<int32_t>(), ElementsAreArray({5, -2, -8, 11}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+TEST(BitwiseXorOpTest, BroadcastRhs) {
+  BitwiseXorOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_UINT32);
+  model.PopulateTensor<uint32_t>(model.input1(), {0, 5, 3, 14});
+  model.PopulateTensor<uint32_t>(model.input2(), {5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput<uint32_t>(), ElementsAreArray({5, 0, 6, 11}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({1, 1, 1, 4}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/broadcast_args.cc

@@ -0,0 +1,114 @@
+/* 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/internal/reference/broadcast_args.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace broadcast_args {
+
+constexpr int kShape1Tensor = 0;
+constexpr int kShape2Tensor = 1;
+constexpr int kOutputTensor = 0;
+
+struct BroadcastArgsContext {
+  BroadcastArgsContext(TfLiteContext* context, TfLiteNode* node) {
+    shape1 = GetInput(context, node, kShape1Tensor);
+    shape2 = GetInput(context, node, kShape2Tensor);
+    output = GetOutput(context, node, kOutputTensor);
+  }
+  const TfLiteTensor* shape1;
+  const TfLiteTensor* shape2;
+  TfLiteTensor* output;
+};
+
+TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node);
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE(context, NumInputs(node) == 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  BroadcastArgsContext op_context(context, node);
+  TF_LITE_ENSURE(context, op_context.shape1->type == kTfLiteInt32 ||
+                              op_context.shape1->type == kTfLiteInt64);
+  TF_LITE_ENSURE_EQ(context, op_context.shape1->type, op_context.shape2->type);
+  TF_LITE_ENSURE_EQ(context, op_context.shape1->type, op_context.output->type);
+
+  // Ensures the shapes are 1D tensor.
+  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.shape1), 1);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.shape2), 1);
+
+  // Resizing the shape of the output tensor.
+  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(1);
+  output_shape->data[0] = std::max(SizeOfDimension(op_context.shape1, 0),
+                                   SizeOfDimension(op_context.shape2, 0));
+  if (IsConstantOrPersistentTensor(op_context.shape1) &&
+      IsConstantOrPersistentTensor(op_context.shape2)) {
+    SetTensorToPersistentRo(op_context.output);
+    TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, op_context.output,
+                                                     output_shape));
+    return EvalImpl(context, node);
+  }
+  return context->ResizeTensor(context, op_context.output, output_shape);
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  BroadcastArgsContext op_context(context, node);
+  if (IsConstantOrPersistentTensor(op_context.output)) {
+    return kTfLiteOk;
+  } else {
+    return EvalImpl(context, node);
+  }
+}
+
+TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) {
+  BroadcastArgsContext op_context(context, node);
+#define TF_LITE_BROADCAST_ARG(data_type)                                    \
+  reference_ops::BroadcastArgs(GetTensorShape(op_context.shape1),           \
+                               GetTensorData<data_type>(op_context.shape1), \
+                               GetTensorShape(op_context.shape2),           \
+                               GetTensorData<data_type>(op_context.shape2), \
+                               GetTensorShape(op_context.output),           \
+                               GetTensorData<data_type>(op_context.output))
+
+  if (op_context.output->type == kTfLiteInt32) {
+    TF_LITE_BROADCAST_ARG(int32_t);
+  } else {
+    TF_LITE_BROADCAST_ARG(int64_t);
+  }
+#undef TF_LITE_BROADCAST_ARG
+
+  return kTfLiteOk;
+}
+
+}  // namespace broadcast_args
+
+TfLiteRegistration* Register_BROADCAST_ARGS() {
+  static TfLiteRegistration r = {nullptr, nullptr, broadcast_args::Prepare,
+                                 broadcast_args::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/broadcast_args_test.cc

@@ -0,0 +1,132 @@
+/* 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 <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace {
+using ::testing::ElementsAreArray;
+
+template <class ShapeType = int32_t>
+class BroadcastArgsOpModel : public SingleOpModel {
+ public:
+  BroadcastArgsOpModel(std::initializer_list<ShapeType> input1,
+                       std::initializer_list<ShapeType> input2,
+                       bool constant_tensor) {
+    int input1_length = input1.size();
+    int input2_length = input2.size();
+    if (constant_tensor) {
+      shape1_ =
+          AddConstInput({GetTensorType<ShapeType>(), {input1_length}}, input1);
+      shape2_ =
+          AddConstInput({GetTensorType<ShapeType>(), {input2_length}}, input2);
+    } else {
+      shape1_ = AddInput({GetTensorType<ShapeType>(), {input1_length}});
+      shape2_ = AddInput({GetTensorType<ShapeType>(), {input2_length}});
+    }
+    output_ = AddOutput(GetTensorType<ShapeType>());
+    SetBuiltinOp(BuiltinOperator_BROADCAST_ARGS, BuiltinOptions_NONE, 0);
+    BuildInterpreter({{input1_length}, {input2_length}});
+    if (!constant_tensor) {
+      if (input1.size() > 0) SetInput1(input1);
+      if (input2.size() > 0) SetInput2(input2);
+    }
+  }
+
+  void SetInput1(std::initializer_list<ShapeType> data) {
+    PopulateTensor(shape1_, data);
+  }
+
+  void SetInput2(std::initializer_list<ShapeType> data) {
+    PopulateTensor(shape2_, data);
+  }
+
+  std::vector<ShapeType> GetOutput() {
+    return ExtractVector<ShapeType>(output_);
+  }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  int shape1_;
+  int shape2_;
+  int output_;
+};
+
+template <typename T>
+class BroadcastArgsOpTest : public ::testing::Test {};
+
+using DataTypes = ::testing::Types<int64_t, int32_t>;
+TYPED_TEST_SUITE(BroadcastArgsOpTest, DataTypes);
+
+#if GTEST_HAS_DEATH_TEST
+TYPED_TEST(BroadcastArgsOpTest, ShapeNotBroadcastableConstant) {
+  EXPECT_DEATH(BroadcastArgsOpModel<TypeParam> m({2, 3, 4, 4}, {2, 2},
+                                                 /*constant_tensor=*/true),
+               "");
+}
+
+TYPED_TEST(BroadcastArgsOpTest, ShapeNotBroadcastable) {
+  BroadcastArgsOpModel<TypeParam> m({2, 3, 4, 4}, {2, 2},
+                                    /*constant_tensor=*/false);
+  EXPECT_DEATH(ASSERT_EQ(m.Invoke(), kTfLiteOk), "");
+}
+#endif
+
+TYPED_TEST(BroadcastArgsOpTest, BroadcastArgsWithScalar) {
+  for (bool constant_tensor : {true, false}) {
+    BroadcastArgsOpModel<TypeParam> m({}, {2, 4}, constant_tensor);
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 4}));
+  }
+}
+
+TYPED_TEST(BroadcastArgsOpTest, BroadcastArgsDifferentDims) {
+  for (bool constant_tensor : {true, false}) {
+    BroadcastArgsOpModel<TypeParam> m({1}, {2, 4}, constant_tensor);
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({2, 4}));
+  }
+}
+
+TYPED_TEST(BroadcastArgsOpTest, BroadcastArgsSameDims) {
+  for (bool constant_tensor : {true, false}) {
+    BroadcastArgsOpModel<TypeParam> m({1, 4, 6, 3, 1, 5}, {4, 4, 1, 3, 4, 1},
+                                      constant_tensor);
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 4, 6, 3, 4, 5}));
+  }
+}
+
+TYPED_TEST(BroadcastArgsOpTest, BroadcastArgsComplex) {
+  for (bool constant_tensor : {true, false}) {
+    BroadcastArgsOpModel<TypeParam> m({6, 3, 1, 5}, {4, 4, 1, 3, 4, 1},
+                                      constant_tensor);
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 4, 6, 3, 4, 5}));
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/broadcast_to.cc

@@ -0,0 +1,128 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/broadcast_to.h"
+
+#include <string.h>
+
+#include <cstdint>
+#include <memory>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace broadcastto {
+
+constexpr int kInputTensor = 0;
+constexpr int kShapeTensor = 1;
+constexpr int kOutputTensor = 0;
+
+struct BroadcastToContext {
+  BroadcastToContext(TfLiteContext* context, TfLiteNode* node) {
+    input = GetInput(context, node, kInputTensor);
+    shape = GetInput(context, node, kShapeTensor);
+    output = GetOutput(context, node, kOutputTensor);
+  }
+  const TfLiteTensor* input;
+  const TfLiteTensor* shape;
+  TfLiteTensor* output;
+};
+
+TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
+                                BroadcastToContext* op_context) {
+  // Ensures the shape is 1D tensor.
+  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->shape), 1);
+
+  // Ensure output dims is not less than input dims.
+  int input_num_dims = NumDimensions(op_context->input);
+  int output_num_dims = SizeOfDimension(op_context->shape, 0);
+  TF_LITE_ENSURE_MSG(context, input_num_dims <= output_num_dims,
+                     "Output shape must be broadcastable from input shape.");
+
+  // Check if output shape is broadcastable from input shape.
+  auto get_shape_data = [op_context](int i) -> int32_t {
+    if (op_context->shape->type == kTfLiteInt32) {
+      return GetTensorData<int32_t>(op_context->shape)[i];
+    } else {
+      return GetTensorData<int64_t>(op_context->shape)[i];
+    }
+  };
+
+  int extending_dims = output_num_dims - input_num_dims;
+  for (int idx = 0; idx < input_num_dims; ++idx) {
+    TF_LITE_ENSURE_MSG(context,
+                       (SizeOfDimension(op_context->input, idx) == 1 ||
+                        SizeOfDimension(op_context->input, idx) ==
+                            get_shape_data(extending_dims + idx)),
+                       "Output shape must be broadcastable from input shape.");
+  }
+  // Resizing the shape of the output tensor.
+  TfLiteIntArray* output_shape = TfLiteIntArrayCreate(output_num_dims);
+  std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>
+      scoped_output_shape(output_shape, TfLiteIntArrayFree);
+  for (int idx = 0; idx < output_num_dims; ++idx) {
+    output_shape->data[idx] = get_shape_data(idx);
+  }
+
+  return context->ResizeTensor(context, op_context->output,
+                               scoped_output_shape.release());
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE(context, NumInputs(node) == 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  BroadcastToContext op_context(context, node);
+  TF_LITE_ENSURE(context, op_context.shape->type == kTfLiteInt32 ||
+                              op_context.shape->type == kTfLiteInt64);
+  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);
+
+  // Not yet support string type due to the use of memcopy with fixed size.
+  TF_LITE_ENSURE(context, op_context.input->type != kTfLiteString);
+
+  if (IsConstantOrPersistentTensor(op_context.shape)) {
+    return ResizeOutputTensor(context, &op_context);
+  }
+
+  SetTensorToDynamic(op_context.output);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  BroadcastToContext op_context(context, node);
+  if (IsDynamicTensor(op_context.output)) {
+    TF_LITE_ENSURE_OK(context, ResizeOutputTensor(context, &op_context));
+  }
+
+  reference_ops::BroadcastTo(
+      GetTensorShape(op_context.input), op_context.input->data.raw,
+      GetTensorShape(op_context.output), op_context.output->data.raw,
+      op_context.input->type);
+  return kTfLiteOk;
+}
+
+}  // namespace broadcastto
+
+TfLiteRegistration* Register_BROADCAST_TO() {
+  static TfLiteRegistration r = {nullptr, nullptr, broadcastto::Prepare,
+                                 broadcastto::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/broadcast_to_test.cc

@@ -0,0 +1,324 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace {
+using ::testing::ElementsAreArray;
+
+template <class InputType, class ShapeType = int32_t>
+class BroadcastToOpModel : public SingleOpModel {
+ public:
+  // BroadcastTo with dynamic shape.
+  BroadcastToOpModel(std::initializer_list<int> input_shape,
+                     std::initializer_list<int> shape_shape) {
+    input_ = AddInput({GetTensorType<InputType>(), input_shape});
+    shape_ = AddInput({GetTensorType<ShapeType>(), shape_shape});
+    output_ = AddOutput(GetTensorType<InputType>());
+    SetBuiltinOp(BuiltinOperator_BROADCAST_TO,
+                 BuiltinOptions_BroadcastToOptions,
+                 CreateBroadcastToOptions(builder_).Union());
+    BuildInterpreter({input_shape, shape_shape});
+  }
+
+  // BroadcastTo with const shape.
+  BroadcastToOpModel(std::initializer_list<int> input_shape,
+                     std::initializer_list<int> shape_shape,
+                     std::initializer_list<ShapeType> shape_values) {
+    input_ = AddInput({GetTensorType<InputType>(), input_shape});
+    shape_ =
+        AddConstInput(GetTensorType<ShapeType>(), shape_values, shape_shape);
+    output_ = AddOutput(GetTensorType<InputType>());
+    SetBuiltinOp(BuiltinOperator_BROADCAST_TO,
+                 BuiltinOptions_BroadcastToOptions,
+                 CreateBroadcastToOptions(builder_).Union());
+    BuildInterpreter({input_shape, shape_shape});
+  }
+
+  void SetInput(std::initializer_list<InputType> data) {
+    PopulateTensor(input_, data);
+  }
+
+  void SetShape(std::initializer_list<ShapeType> data) {
+    PopulateTensor(shape_, data);
+  }
+
+  std::vector<InputType> GetOutput() {
+    return ExtractVector<InputType>(output_);
+  }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ protected:
+  int input_;
+  int shape_;
+  int output_;
+};
+
+template <typename T>
+class BroadcastToOpTest : public ::testing::Test {};
+
+using DataTypes = ::testing::Types<float, uint8_t, int8_t, int16_t, int32_t>;
+TYPED_TEST_SUITE(BroadcastToOpTest, DataTypes);
+
+#if GTEST_HAS_DEATH_TEST
+TYPED_TEST(BroadcastToOpTest, ShapeMustBe1D) {
+  EXPECT_DEATH(
+      BroadcastToOpModel<TypeParam>({2, 3, 4, 4}, {2, 2}, {2, 3, 4, 4}), "");
+  // Non-constant Shape tensor.
+  BroadcastToOpModel<TypeParam> m({2, 3, 4, 4}, {2, 2});
+  m.SetShape({2, 3, 4, 4});
+  EXPECT_THAT(m.Invoke(), kTfLiteError);
+}
+
+TYPED_TEST(BroadcastToOpTest, MismatchDimension) {
+  EXPECT_DEATH(BroadcastToOpModel<TypeParam>({2, 4, 1, 2}, {4}, {2, 4, 1, 3}),
+               "Output shape must be broadcastable from input shape.");
+  EXPECT_DEATH(
+      BroadcastToOpModel<TypeParam>({2, 4, 1, 2, 3}, {4}, {2, 4, 1, 2}),
+      "Output shape must be broadcastable from input shape.");
+
+  // Non-constant Shape tensor.
+  BroadcastToOpModel<TypeParam> m1({2, 4, 1, 2}, {4});
+  m1.SetShape({2, 3, 4, 4});
+  EXPECT_THAT(m1.Invoke(), kTfLiteError);
+  BroadcastToOpModel<TypeParam> m2({2, 4, 1, 2}, {5});
+  m2.SetShape({1, 2, 3, 4, 4});
+  EXPECT_THAT(m2.Invoke(), kTfLiteError);
+}
+#endif
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo1DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1}, {1}, {4});
+  m.SetInput({3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 3, 3}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo4DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 2}, {4}, {1, 1, 2, 2});
+  m.SetInput({3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 4, 3, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo8DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 1, 1, 1, 2, 1}, {8},
+                                  {1, 1, 1, 1, 1, 1, 2, 2});
+  m.SetInput({3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo9DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 1, 1, 1, 1, 2, 1}, {9},
+                                  {1, 1, 1, 1, 1, 1, 1, 2, 2});
+  m.SetInput({3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(),
+              ElementsAreArray({1, 1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo1DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1}, {1});
+  m.SetInput({3});
+  m.SetShape({4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 3, 3}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo4DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 2}, {4});
+  m.SetInput({3, 4});
+  m.SetShape({1, 1, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 4, 3, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo8DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 1, 1, 1, 2, 1}, {8});
+  m.SetInput({3, 4});
+  m.SetShape({1, 1, 1, 1, 1, 1, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastTo9DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1, 1, 1, 1, 1, 1, 1, 2, 1}, {9});
+  m.SetInput({3, 4});
+  m.SetShape({1, 1, 1, 1, 1, 1, 1, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(),
+              ElementsAreArray({1, 1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast4DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 3, 1, 2}, {4}, {3, 3, 2, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 3, 2, 2}));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4,
+                        3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast4DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1, 3, 1, 2}, {4});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  m.SetShape({3, 3, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 3, 2, 2}));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4,
+                        3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast6DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 2, 1, 3, 1, 2}, {6}, {2, 2, 1, 3, 2, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 1, 3, 2, 2}));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({1, 2, 1, 2, 3, 4,  3, 4,  5,  6,  5,  6,
+                                7, 8, 7, 8, 9, 10, 9, 10, 11, 12, 11, 12,
+                                1, 2, 1, 2, 3, 4,  3, 4,  5,  6,  5,  6,
+                                7, 8, 7, 8, 9, 10, 9, 10, 11, 12, 11, 12}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast6DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({1, 2, 1, 3, 1, 2}, {6});
+  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  m.SetShape({2, 2, 1, 3, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 1, 3, 2, 2}));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({1, 2, 1, 2, 3, 4,  3, 4,  5,  6,  5,  6,
+                                7, 8, 7, 8, 9, 10, 9, 10, 11, 12, 11, 12,
+                                1, 2, 1, 2, 3, 4,  3, 4,  5,  6,  5,  6,
+                                7, 8, 7, 8, 9, 10, 9, 10, 11, 12, 11, 12}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast8DConstTest) {
+  BroadcastToOpModel<TypeParam> m({1, 3, 1, 2, 1, 4, 1, 1}, {8},
+                                  {2, 3, 1, 2, 2, 4, 1, 1});
+  m.SetInput({1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+              13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3, 1, 2, 2, 4, 1, 1}));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1,  2,  3,  4,  1,  2,  3,  4,  5,  6,  7,  8,  5,  6,
+                        7,  8,  9,  10, 11, 12, 9,  10, 11, 12, 13, 14, 15, 16,
+                        13, 14, 15, 16, 17, 18, 19, 20, 17, 18, 19, 20, 21, 22,
+                        23, 24, 21, 22, 23, 24, 1,  2,  3,  4,  1,  2,  3,  4,
+                        5,  6,  7,  8,  5,  6,  7,  8,  9,  10, 11, 12, 9,  10,
+                        11, 12, 13, 14, 15, 16, 13, 14, 15, 16, 17, 18, 19, 20,
+                        17, 18, 19, 20, 21, 22, 23, 24, 21, 22, 23, 24}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ComplexBroadcast8DDynamicTest) {
+  BroadcastToOpModel<TypeParam> m({2, 1, 1, 2, 1, 4, 1, 1}, {8});
+  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  m.SetShape({2, 3, 2, 2, 2, 4, 1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 3, 2, 2, 2, 4, 1, 1}));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           1, 2,  3,  4,  1, 2,  3,  4,  5,  6,  7,  8,  5,  6,  7,  8,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16,
+           9, 10, 11, 12, 9, 10, 11, 12, 13, 14, 15, 16, 13, 14, 15, 16}));
+}
+
+TYPED_TEST(BroadcastToOpTest, ExtendingShape4DConstTest) {
+  BroadcastToOpModel<TypeParam> m({3, 1, 2}, {4}, {3, 3, 2, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 3, 2, 2}));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4,
+                        3, 4, 5, 6, 5, 6, 1, 2, 1, 2, 3, 4, 3, 4, 5, 6, 5, 6}));
+}
+
+TYPED_TEST(BroadcastToOpTest, NoBroadcastingConstTest) {
+  BroadcastToOpModel<TypeParam> m({3, 1, 2}, {3}, {3, 1, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 1, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TYPED_TEST(BroadcastToOpTest, NoBroadcasting8DConstTest) {
+  BroadcastToOpModel<TypeParam> m({3, 1, 1, 1, 1, 1, 1, 2}, {8},
+                                  {3, 1, 1, 1, 1, 1, 1, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 1, 1, 1, 1, 1, 1, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TYPED_TEST(BroadcastToOpTest, Int64ShapeConstTest) {
+  BroadcastToOpModel<TypeParam, int64_t> m({1, 1, 1, 1, 1, 1, 2, 1}, {8},
+                                           {1, 1, 1, 1, 1, 1, 2, 2});
+  m.SetInput({3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, Int64ShapeDDynamicTest) {
+  BroadcastToOpModel<TypeParam, int64_t> m({1, 1, 1, 1, 1, 1, 2, 1}, {8});
+  m.SetInput({3, 4});
+  m.SetShape({1, 1, 1, 1, 1, 1, 2, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 1, 2, 2}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 4, 4}));
+}
+
+TYPED_TEST(BroadcastToOpTest, BroadcastToEmtpyShapeTest) {
+  BroadcastToOpModel<TypeParam> m({3, 1, 2}, {3}, {3, 0, 2});
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3, 0, 2}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/bucketize.cc

@@ -0,0 +1,157 @@
+/* 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 <stdint.h>
+
+#include <algorithm>
+
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace bucketize {
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+struct OpData {
+  // boundaries array is owned by the buffer housing TfLiteBucketizeParams.
+  const float* boundaries;
+  int num_boundaries;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  const auto* params = reinterpret_cast<const TfLiteBucketizeParams*>(buffer);
+
+  if (!FLATBUFFERS_LITTLEENDIAN) {
+    int32_t* p =
+        reinterpret_cast<int32_t*>(const_cast<float*>(params->boundaries));
+    for (size_t i = 0; i < params->num_boundaries; i++, p++)
+      *p = flatbuffers::EndianSwap(*p);
+  }
+
+  op_data->boundaries = params->boundaries;
+  op_data->num_boundaries = params->num_boundaries;
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+  if (!std::is_sorted(opdata->boundaries,
+                      opdata->boundaries + opdata->num_boundaries)) {
+    TF_LITE_KERNEL_LOG(context, "Expected sorted boundaries");
+    return kTfLiteError;
+  }
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+
+  if (input->type != kTfLiteInt32 && input->type != kTfLiteFloat32 &&
+      input->type != kTfLiteInt64 && input->type != kTfLiteFloat64) {
+    TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by bucketize.",
+                       TfLiteTypeGetName(input->type));
+    return kTfLiteError;
+  }
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  output->type = kTfLiteInt32;
+
+  TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
+  return context->ResizeTensor(context, output, output_shape);
+}
+
+template <typename T>
+inline void Bucketize(const RuntimeShape& input_shape, const T* input_data,
+                      const float* boundaries, int num_boundaries,
+                      const RuntimeShape& output_shape, int32_t* output_data) {
+  const int flat_size = MatchingFlatSize(input_shape, output_shape);
+
+  for (int i = 0; i < flat_size; i++) {
+    auto first_bigger_it = std::upper_bound(
+        boundaries, boundaries + num_boundaries, input_data[i]);
+    output_data[i] = first_bigger_it - boundaries;
+  }
+}
+
+template <typename T>
+TfLiteStatus BucketizeImpl(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt32);
+
+  Bucketize<T>(GetTensorShape(input), GetTensorData<T>(input),
+               opdata->boundaries, opdata->num_boundaries,
+               GetTensorShape(output), GetTensorData<int32_t>(output));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+
+  switch (input->type) {
+    case kTfLiteFloat32: {
+      return BucketizeImpl<float>(context, node);
+    }
+    case kTfLiteFloat64: {
+      return BucketizeImpl<double>(context, node);
+    }
+    case kTfLiteInt32: {
+      return BucketizeImpl<int32_t>(context, node);
+    }
+    case kTfLiteInt64: {
+      return BucketizeImpl<int64_t>(context, node);
+    }
+    default: {
+      TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported by bucketize.",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+    }
+  }
+}
+
+}  // namespace
+}  // namespace bucketize
+
+TfLiteRegistration* Register_BUCKETIZE() {
+  static TfLiteRegistration r = {bucketize::Init, bucketize::Free,
+                                 bucketize::Prepare, bucketize::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/bucketize_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 <stdint.h>
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+template <typename T>
+class BucketizeOpModel : public SingleOpModel {
+ public:
+  BucketizeOpModel(const TensorData& input,
+                   const std::vector<float>& boundaries) {
+    input_ = AddInput(input);
+    boundaries_ = boundaries;
+
+    output_ = AddOutput({TensorType_INT32, input.shape});
+
+    SetBuiltinOp(BuiltinOperator_BUCKETIZE, BuiltinOptions_BucketizeOptions,
+                 CreateBucketizeOptions(
+                     builder_, builder_.CreateVector<float>(boundaries_))
+                     .Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input() { return input_; }
+  const std::vector<float>& boundaries() { return boundaries_; }
+
+  std::vector<int> GetOutput() { return ExtractVector<int>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input_;
+  std::vector<float> boundaries_;
+  int output_;
+};
+
+TEST(BucketizeOpTest, Float) {
+  // Buckets are: (-inf, 0.), [0., 10.), [10., 100.), [100., +inf).
+  BucketizeOpModel<float> model(
+      /*input=*/{/*type=*/TensorType_FLOAT32, /*shape=*/{3, 2}},
+      /*boundaries=*/{0.0f, 10.0f, 100.0f});
+
+  // input: [[-5, 10000], [150, 10], [5, 100]]
+  model.PopulateTensor<float>(model.input(),
+                              {-5.0f, 10000.0f, 150.0f, 10.0f, 5.0f, 100.0f});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  // output: [[0, 3], [3, 2], [1, 3]]
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(3, 2));
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({0, 3, 3, 2, 1, 3}));
+}
+
+TEST(BucketizeOpTest, Int32) {
+  // Buckets are: (-inf, 0.), [0., 10.), [10., 100.), [100., +inf).
+  BucketizeOpModel<int32_t> model(
+      /*input=*/{/*type=*/TensorType_INT32, /*shape=*/{3, 2}},
+      /*boundaries=*/{0, 10, 100});
+
+  // input: [[-5, 10000], [150, 10], [5, 100]]
+  model.PopulateTensor<int32_t>(model.input(), {-5, 10000, 150, 10, 5, 100});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  // output: [[0, 3], [3, 2], [1, 3]]
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(3, 2));
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({0, 3, 3, 2, 1, 3}));
+}
+
+#if GTEST_HAS_DEATH_TEST
+TEST(BucketizeOpTest, UnsortedBuckets) {
+  EXPECT_DEATH(BucketizeOpModel<float>(
+                   /*input=*/{/*type=*/TensorType_INT32, /*shape=*/{3, 2}},
+                   /*boundaries=*/{0, 10, -10}),
+               "Expected sorted boundaries");
+}
+#endif
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/builtin_op_kernels.h

@@ -0,0 +1,38 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_BUILTIN_OP_KERNELS_H_
+#define TENSORFLOW_LITE_KERNELS_BUILTIN_OP_KERNELS_H_
+
+/// For documentation, see
+/// third_party/tensorflow/lite/core/kernels/builtin_op_kernels.h
+
+#include "tensorflow/lite/core/kernels/builtin_op_kernels.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+
+#define TFLITE_OP(NAME) \
+    using ::tflite::ops::builtin::NAME;
+
+#include "tensorflow/lite/kernels/builtin_ops_list.inc"
+
+#undef TFLITE_OP
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_BUILTIN_OP_KERNELS_H_

tensorflow/lite/kernels/builtin_ops_list.inc

@@ -0,0 +1,224 @@
+
+/* 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.
+==============================================================================*/
+
+// DO NOT EDIT MANUALLY: This file is automatically generated by
+// `tensorflow/lite/schema/builtin_ops_list/generator.cc`.
+
+TFLITE_OP(Register_ADD)
+TFLITE_OP(Register_AVERAGE_POOL_2D)
+TFLITE_OP(Register_CONCATENATION)
+TFLITE_OP(Register_CONV_2D)
+TFLITE_OP(Register_DEPTHWISE_CONV_2D)
+TFLITE_OP(Register_DEPTH_TO_SPACE)
+TFLITE_OP(Register_DEQUANTIZE)
+TFLITE_OP(Register_EMBEDDING_LOOKUP)
+TFLITE_OP(Register_FLOOR)
+TFLITE_OP(Register_FULLY_CONNECTED)
+TFLITE_OP(Register_HASHTABLE_LOOKUP)
+TFLITE_OP(Register_L2_NORMALIZATION)
+TFLITE_OP(Register_L2_POOL_2D)
+TFLITE_OP(Register_LOCAL_RESPONSE_NORMALIZATION)
+TFLITE_OP(Register_LOGISTIC)
+TFLITE_OP(Register_LSH_PROJECTION)
+TFLITE_OP(Register_LSTM)
+TFLITE_OP(Register_MAX_POOL_2D)
+TFLITE_OP(Register_MUL)
+TFLITE_OP(Register_RELU)
+TFLITE_OP(Register_RELU_N1_TO_1)
+TFLITE_OP(Register_RELU6)
+TFLITE_OP(Register_RESHAPE)
+TFLITE_OP(Register_RESIZE_BILINEAR)
+TFLITE_OP(Register_RNN)
+TFLITE_OP(Register_SOFTMAX)
+TFLITE_OP(Register_SPACE_TO_DEPTH)
+TFLITE_OP(Register_SVDF)
+TFLITE_OP(Register_TANH)
+TFLITE_OP(Register_SKIP_GRAM)
+TFLITE_OP(Register_EMBEDDING_LOOKUP_SPARSE)
+TFLITE_OP(Register_PAD)
+TFLITE_OP(Register_UNIDIRECTIONAL_SEQUENCE_RNN)
+TFLITE_OP(Register_GATHER)
+TFLITE_OP(Register_BATCH_TO_SPACE_ND)
+TFLITE_OP(Register_SPACE_TO_BATCH_ND)
+TFLITE_OP(Register_TRANSPOSE)
+TFLITE_OP(Register_MEAN)
+TFLITE_OP(Register_SUB)
+TFLITE_OP(Register_DIV)
+TFLITE_OP(Register_SQUEEZE)
+TFLITE_OP(Register_UNIDIRECTIONAL_SEQUENCE_LSTM)
+TFLITE_OP(Register_STRIDED_SLICE)
+TFLITE_OP(Register_BIDIRECTIONAL_SEQUENCE_RNN)
+TFLITE_OP(Register_EXP)
+TFLITE_OP(Register_TOPK_V2)
+TFLITE_OP(Register_SPLIT)
+TFLITE_OP(Register_LOG_SOFTMAX)
+TFLITE_OP(Register_BIDIRECTIONAL_SEQUENCE_LSTM)
+TFLITE_OP(Register_CAST)
+TFLITE_OP(Register_PRELU)
+TFLITE_OP(Register_MAXIMUM)
+TFLITE_OP(Register_ARG_MAX)
+TFLITE_OP(Register_MINIMUM)
+TFLITE_OP(Register_LESS)
+TFLITE_OP(Register_NEG)
+TFLITE_OP(Register_PADV2)
+TFLITE_OP(Register_GREATER)
+TFLITE_OP(Register_GREATER_EQUAL)
+TFLITE_OP(Register_LESS_EQUAL)
+TFLITE_OP(Register_SELECT)
+TFLITE_OP(Register_SLICE)
+TFLITE_OP(Register_SIN)
+TFLITE_OP(Register_TRANSPOSE_CONV)
+TFLITE_OP(Register_SPARSE_TO_DENSE)
+TFLITE_OP(Register_TILE)
+TFLITE_OP(Register_EXPAND_DIMS)
+TFLITE_OP(Register_EQUAL)
+TFLITE_OP(Register_NOT_EQUAL)
+TFLITE_OP(Register_LOG)
+TFLITE_OP(Register_SUM)
+TFLITE_OP(Register_SQRT)
+TFLITE_OP(Register_RSQRT)
+TFLITE_OP(Register_SHAPE)
+TFLITE_OP(Register_POW)
+TFLITE_OP(Register_ARG_MIN)
+TFLITE_OP(Register_FAKE_QUANT)
+TFLITE_OP(Register_REDUCE_PROD)
+TFLITE_OP(Register_REDUCE_MAX)
+TFLITE_OP(Register_PACK)
+TFLITE_OP(Register_LOGICAL_OR)
+TFLITE_OP(Register_ONE_HOT)
+TFLITE_OP(Register_LOGICAL_AND)
+TFLITE_OP(Register_LOGICAL_NOT)
+TFLITE_OP(Register_UNPACK)
+TFLITE_OP(Register_REDUCE_MIN)
+TFLITE_OP(Register_FLOOR_DIV)
+TFLITE_OP(Register_REDUCE_ANY)
+TFLITE_OP(Register_SQUARE)
+TFLITE_OP(Register_ZEROS_LIKE)
+TFLITE_OP(Register_FILL)
+TFLITE_OP(Register_FLOOR_MOD)
+TFLITE_OP(Register_RANGE)
+TFLITE_OP(Register_RESIZE_NEAREST_NEIGHBOR)
+TFLITE_OP(Register_LEAKY_RELU)
+TFLITE_OP(Register_SQUARED_DIFFERENCE)
+TFLITE_OP(Register_MIRROR_PAD)
+TFLITE_OP(Register_ABS)
+TFLITE_OP(Register_SPLIT_V)
+TFLITE_OP(Register_UNIQUE)
+TFLITE_OP(Register_CEIL)
+TFLITE_OP(Register_REVERSE_V2)
+TFLITE_OP(Register_ADD_N)
+TFLITE_OP(Register_GATHER_ND)
+TFLITE_OP(Register_COS)
+TFLITE_OP(Register_WHERE)
+TFLITE_OP(Register_RANK)
+TFLITE_OP(Register_ELU)
+TFLITE_OP(Register_REVERSE_SEQUENCE)
+TFLITE_OP(Register_MATRIX_DIAG)
+TFLITE_OP(Register_QUANTIZE)
+TFLITE_OP(Register_MATRIX_SET_DIAG)
+TFLITE_OP(Register_ROUND)
+TFLITE_OP(Register_HARD_SWISH)
+TFLITE_OP(Register_IF)
+TFLITE_OP(Register_WHILE)
+TFLITE_OP(Register_NON_MAX_SUPPRESSION_V4)
+TFLITE_OP(Register_NON_MAX_SUPPRESSION_V5)
+TFLITE_OP(Register_SCATTER_ND)
+TFLITE_OP(Register_SELECT_V2)
+TFLITE_OP(Register_DENSIFY)
+TFLITE_OP(Register_SEGMENT_SUM)
+TFLITE_OP(Register_BATCH_MATMUL)
+TFLITE_OP(Register_CUMSUM)
+TFLITE_OP(Register_CALL_ONCE)
+TFLITE_OP(Register_BROADCAST_TO)
+TFLITE_OP(Register_RFFT2D)
+TFLITE_OP(Register_CONV_3D)
+TFLITE_OP(Register_IMAG)
+TFLITE_OP(Register_REAL)
+TFLITE_OP(Register_COMPLEX_ABS)
+TFLITE_OP(Register_HASHTABLE)
+TFLITE_OP(Register_HASHTABLE_FIND)
+TFLITE_OP(Register_HASHTABLE_IMPORT)
+TFLITE_OP(Register_HASHTABLE_SIZE)
+TFLITE_OP(Register_REDUCE_ALL)
+TFLITE_OP(Register_CONV_3D_TRANSPOSE)
+TFLITE_OP(Register_VAR_HANDLE)
+TFLITE_OP(Register_READ_VARIABLE)
+TFLITE_OP(Register_ASSIGN_VARIABLE)
+TFLITE_OP(Register_BROADCAST_ARGS)
+TFLITE_OP(Register_RANDOM_STANDARD_NORMAL)
+TFLITE_OP(Register_BUCKETIZE)
+TFLITE_OP(Register_RANDOM_UNIFORM)
+TFLITE_OP(Register_MULTINOMIAL)
+TFLITE_OP(Register_GELU)
+TFLITE_OP(Register_DYNAMIC_UPDATE_SLICE)
+TFLITE_OP(Register_RELU_0_TO_1)
+TFLITE_OP(Register_UNSORTED_SEGMENT_PROD)
+TFLITE_OP(Register_UNSORTED_SEGMENT_MAX)
+TFLITE_OP(Register_UNSORTED_SEGMENT_SUM)
+TFLITE_OP(Register_ATAN2)
+TFLITE_OP(Register_UNSORTED_SEGMENT_MIN)
+TFLITE_OP(Register_SIGN)
+TFLITE_OP(Register_BITCAST)
+TFLITE_OP(Register_BITWISE_XOR)
+TFLITE_OP(Register_RIGHT_SHIFT)
+TFLITE_OP(Register_STABLEHLO_LOGISTIC)
+TFLITE_OP(Register_STABLEHLO_ADD)
+TFLITE_OP(Register_STABLEHLO_DIVIDE)
+TFLITE_OP(Register_STABLEHLO_MULTIPLY)
+TFLITE_OP(Register_STABLEHLO_MAXIMUM)
+TFLITE_OP(Register_STABLEHLO_RESHAPE)
+TFLITE_OP(Register_STABLEHLO_CLAMP)
+TFLITE_OP(Register_STABLEHLO_CONCATENATE)
+TFLITE_OP(Register_STABLEHLO_BROADCAST_IN_DIM)
+TFLITE_OP(Register_STABLEHLO_CONVOLUTION)
+TFLITE_OP(Register_STABLEHLO_SLICE)
+TFLITE_OP(Register_STABLEHLO_CUSTOM_CALL)
+TFLITE_OP(Register_STABLEHLO_REDUCE)
+TFLITE_OP(Register_STABLEHLO_ABS)
+TFLITE_OP(Register_STABLEHLO_AND)
+TFLITE_OP(Register_STABLEHLO_COSINE)
+TFLITE_OP(Register_STABLEHLO_EXPONENTIAL)
+TFLITE_OP(Register_STABLEHLO_FLOOR)
+TFLITE_OP(Register_STABLEHLO_LOG)
+TFLITE_OP(Register_STABLEHLO_MINIMUM)
+TFLITE_OP(Register_STABLEHLO_NEGATE)
+TFLITE_OP(Register_STABLEHLO_OR)
+TFLITE_OP(Register_STABLEHLO_POWER)
+TFLITE_OP(Register_STABLEHLO_REMAINDER)
+TFLITE_OP(Register_STABLEHLO_RSQRT)
+TFLITE_OP(Register_STABLEHLO_SELECT)
+TFLITE_OP(Register_STABLEHLO_SUBTRACT)
+TFLITE_OP(Register_STABLEHLO_TANH)
+TFLITE_OP(Register_STABLEHLO_SCATTER)
+TFLITE_OP(Register_STABLEHLO_COMPARE)
+TFLITE_OP(Register_STABLEHLO_CONVERT)
+TFLITE_OP(Register_STABLEHLO_DYNAMIC_SLICE)
+TFLITE_OP(Register_STABLEHLO_DYNAMIC_UPDATE_SLICE)
+TFLITE_OP(Register_STABLEHLO_PAD)
+TFLITE_OP(Register_STABLEHLO_IOTA)
+TFLITE_OP(Register_STABLEHLO_DOT_GENERAL)
+TFLITE_OP(Register_STABLEHLO_REDUCE_WINDOW)
+TFLITE_OP(Register_STABLEHLO_SORT)
+TFLITE_OP(Register_STABLEHLO_WHILE)
+TFLITE_OP(Register_STABLEHLO_GATHER)
+TFLITE_OP(Register_STABLEHLO_TRANSPOSE)
+TFLITE_OP(Register_DILATE)
+TFLITE_OP(Register_STABLEHLO_RNG_BIT_GENERATOR)
+TFLITE_OP(Register_REDUCE_WINDOW)
+TFLITE_OP(Register_STABLEHLO_COMPOSITE)
+TFLITE_OP(Register_STABLEHLO_SHIFT_LEFT)
+TFLITE_OP(Register_STABLEHLO_CBRT)
+TFLITE_OP(Register_STABLEHLO_CASE)

tensorflow/lite/kernels/call_once.cc

@@ -0,0 +1,122 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <stddef.h>
+
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/experimental/resource/initialization_status.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace call_once_kernel {
+
+// CallOnce operator is a control flow op to invoke other subgraph in the graph
+// in order to conduct the given graph's initialization tasks, for example, hash
+// table initialization and variable initialization.
+//
+// This operator will invoke the subgraph for initialization in the first run
+// and become no-op after the first run in an interpreter's life cycle.
+
+struct OpData {
+  // Subgraph index to be invoked once in a life cycle by this CallOnce op.
+  int init_subgraph_index;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData;
+  const auto* params = reinterpret_cast<const TfLiteCallOnceParams*>(buffer);
+  op_data->init_subgraph_index = params->init_subgraph_index;
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  const OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+
+  // The CALL_ONCE op cannot call the same subgraph recursively or it will cause
+  // an infinite recursion.
+  TF_LITE_ENSURE(context, op_data->init_subgraph_index !=
+                              this_subgraph->GetSubgraphIndex());
+
+  // Return early if the initialization graph is already invoked.
+  resource::InitializationStatusMap* map =
+      &this_subgraph->initialization_status_map();
+  resource::InitializationStatus* status =
+      resource::GetInitializationStatus(map, op_data->init_subgraph_index);
+  if (status->IsInitialized()) return kTfLiteOk;
+
+  auto* subgraphs = this_subgraph->GetSubgraphs();
+
+  TF_LITE_ENSURE_EQ(context, node->inputs->size, 0);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 0);
+
+  TF_LITE_ENSURE(context, op_data->init_subgraph_index < subgraphs->size());
+
+  // Ensures that there are no input and output tensors in the subgraph.
+  Subgraph* init_subgraph = (*subgraphs)[op_data->init_subgraph_index].get();
+  TF_LITE_ENSURE_EQ(context, init_subgraph->inputs().size(), 0);
+  TF_LITE_ENSURE_EQ(context, init_subgraph->outputs().size(), 0);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+
+  Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+
+  // The initialization graph should be invoked once in a life cycle.
+  resource::InitializationStatusMap* map =
+      &this_subgraph->initialization_status_map();
+  resource::InitializationStatus* status =
+      resource::GetInitializationStatus(map, op_data->init_subgraph_index);
+  if (status->IsInitialized()) return kTfLiteOk;
+
+  auto* subgraphs = this_subgraph->GetSubgraphs();
+  Subgraph& init_subgraph = *(*subgraphs)[op_data->init_subgraph_index];
+
+  TF_LITE_ENSURE_OK(context, init_subgraph.AllocateTensors());
+  TF_LITE_ENSURE_OK(context, init_subgraph.Invoke());
+  TF_LITE_ENSURE_OK(context, init_subgraph.ReleaseNonPersistentMemory());
+
+  // Mark the invocation completed.
+  status->MarkInitializationIsDone();
+  return kTfLiteOk;
+}
+
+}  // namespace call_once_kernel
+
+TfLiteRegistration* Register_CALL_ONCE() {
+  static TfLiteRegistration r = {call_once_kernel::Init, call_once_kernel::Free,
+                                 call_once_kernel::Prepare,
+                                 call_once_kernel::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/call_once_test.cc

@@ -0,0 +1,116 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <stdint.h>
+
+#include <memory>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/subgraph_test_util.h"
+
+namespace tflite {
+
+using subgraph_test_util::ControlFlowOpTest;
+
+namespace {
+
+class CallOnceTest : public ControlFlowOpTest {
+ protected:
+  void SetUp() override {
+    AddSubgraphs(2);
+    builder_->BuildCallOnceAndReadVariableSubgraph(
+        &interpreter_->primary_subgraph());
+    builder_->BuildAssignRandomValueToVariableSubgraph(
+        interpreter_->subgraph(1));
+    builder_->BuildCallOnceAndReadVariablePlusOneSubgraph(
+        interpreter_->subgraph(2));
+
+    ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+    ASSERT_EQ(interpreter_->subgraph(2)->AllocateTensors(), kTfLiteOk);
+  }
+};
+
+TEST_F(CallOnceTest, TestSimple) {
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  ASSERT_EQ(output->dims->size, 1);
+  ASSERT_EQ(output->dims->data[0], 1);
+  ASSERT_EQ(output->type, kTfLiteInt32);
+  ASSERT_EQ(NumElements(output), 1);
+
+  // The value of the variable must be non-zero, which will be assigned by the
+  // initialization subgraph.
+  EXPECT_GT(output->data.i32[0], 0);
+}
+
+TEST_F(CallOnceTest, TestInvokeMultipleTimes) {
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  ASSERT_EQ(output->dims->size, 1);
+  ASSERT_EQ(output->dims->data[0], 1);
+  ASSERT_EQ(output->type, kTfLiteInt32);
+  ASSERT_EQ(NumElements(output), 1);
+
+  // The value of the variable must be non-zero, which will be assigned by the
+  // initialization subgraph.
+  int value = output->data.i32[0];
+  EXPECT_GT(value, 0);
+
+  for (int i = 0; i < 3; ++i) {
+    // Make sure that no more random value assignment in the initialization
+    // subgraph.
+    ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+    TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+    ASSERT_EQ(output->dims->size, 1);
+    ASSERT_EQ(output->dims->data[0], 1);
+    ASSERT_EQ(output->type, kTfLiteInt32);
+    ASSERT_EQ(NumElements(output), 1);
+    ASSERT_EQ(output->data.i32[0], value);
+  }
+}
+
+TEST_F(CallOnceTest, TestInvokeOnceAcrossMultipleEntryPoints) {
+  ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+
+  TfLiteTensor* output = interpreter_->tensor(interpreter_->outputs()[0]);
+  ASSERT_EQ(output->dims->size, 1);
+  ASSERT_EQ(output->dims->data[0], 1);
+  ASSERT_EQ(output->type, kTfLiteInt32);
+  ASSERT_EQ(NumElements(output), 1);
+
+  // The value of the variable must be non-zero, which will be assigned by the
+  // initialization subgraph.
+  int value = output->data.i32[0];
+  EXPECT_GT(value, 0);
+
+  // Make sure that no more random value assignment in the initialization
+  // subgraph while invoking the other subgraph, which has the CallOnce op.
+  ASSERT_EQ(interpreter_->subgraph(2)->Invoke(), kTfLiteOk);
+  output = interpreter_->subgraph(2)->tensor(
+      interpreter_->subgraph(2)->outputs()[0]);
+  ASSERT_EQ(output->dims->size, 1);
+  ASSERT_EQ(output->dims->data[0], 1);
+  ASSERT_EQ(output->type, kTfLiteInt32);
+  ASSERT_EQ(NumElements(output), 1);
+  ASSERT_EQ(output->data.i32[0], value + 1);
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/cast.cc

@@ -0,0 +1,522 @@
+/* Copyright 2017 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 <algorithm>
+#include <cmath>
+#include <complex>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/interpreter_options.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/types/fp16.h"
+#include "tensorflow/lite/types/half.h"
+
+#ifdef __ARM_NEON
+#include <arm_neon.h>
+#endif
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace cast {
+
+namespace {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+void copyCast(const float* in, int32_t* out, int num_elements) {
+  const float min_int_float =
+      static_cast<float>(std::numeric_limits<int32_t>::min());
+  const float max_int_float = std::nextafterf(
+      static_cast<float>(std::numeric_limits<int32_t>::max()), 0);
+
+  std::transform(in, in + num_elements, out, [=](float a) {
+    return a <= max_int_float ? static_cast<int32_t>(std::max(a, min_int_float))
+                              : std::numeric_limits<int32_t>::max();
+  });
+}
+
+void copyCast(const float* in, int16_t* out, int num_elements) {
+  const float min_int_float =
+      static_cast<float>(std::numeric_limits<int16_t>::min());
+  const float max_int_float =
+      static_cast<float>(std::numeric_limits<int16_t>::max());
+  std::transform(in, in + num_elements, out, [=](float a) {
+    return static_cast<int16_t>(
+        std::max(std::min(a, max_int_float), min_int_float));
+  });
+}
+
+void copyCast(const float* in, uint8_t* out, int num_elements) {
+  const float min_int_float =
+      static_cast<float>(std::numeric_limits<uint8_t>::min());
+  const float max_int_float =
+      static_cast<float>(std::numeric_limits<uint8_t>::max());
+  std::transform(in, in + num_elements, out, [=](float a) {
+    return static_cast<uint8_t>(
+        std::max(std::min(a, max_int_float), min_int_float));
+  });
+}
+
+template <typename FromT, typename ToT>
+void copyCast(const FromT* in, ToT* out, int num_elements) {
+  std::transform(in, in + num_elements, out,
+                 [](FromT a) { return static_cast<ToT>(a); });
+}
+
+template <typename ToT>
+void copyCast(const std::complex<float>* in, ToT* out, int num_elements) {
+  std::transform(in, in + num_elements, out, [](std::complex<float> a) {
+    return static_cast<ToT>(std::real(a));
+  });
+}
+
+template <>
+void copyCast(const std::complex<float>* in, std::complex<float>* out,
+              int num_elements) {
+  std::transform(in, in + num_elements, out,
+                 [](std::complex<float> a) { return a; });
+}
+
+template <typename ToT>
+void copyCast(const half* in, ToT* out, int num_elements) {
+  std::transform(in, in + num_elements, out, [](half a) {
+    return static_cast<ToT>(fp16_ieee_to_fp32_value(a));
+  });
+}
+
+template <>
+void copyCast(const Eigen::bfloat16* in, std::complex<float>* out,
+              int num_elements) {
+  std::transform(in, in + num_elements, out, [](Eigen::bfloat16 a) {
+    return std::complex<float>(Eigen::bfloat16_impl::bfloat16_to_float(a));
+  });
+}
+
+template <typename FromT>
+void copyCastToFloat16(const FromT* in, half* out, int num_elements) {
+  std::transform(in, in + num_elements, out, [](FromT a) {
+    return half::from_bits(fp16_ieee_from_fp32_value(static_cast<float>(a)));
+  });
+}
+
+template <>
+void copyCastToFloat16(const std::complex<float>* in, half* out,
+                       int num_elements) {
+  std::transform(in, in + num_elements, out, [](std::complex<float> a) {
+    return half::from_bits(fp16_ieee_from_fp32_value(std::real(a)));
+  });
+}
+
+template <>
+void copyCastToFloat16(const Eigen::bfloat16* in, half* out, int num_elements) {
+  // bfloat16 -> float -> half (fp16)
+  std::transform(in, in + num_elements, out, [](Eigen::bfloat16 a) {
+    return half::from_bits(
+        fp16_ieee_from_fp32_value(Eigen::bfloat16_impl::bfloat16_to_float(a)));
+  });
+}
+
+template <typename FromT>
+void copyCastToBFloat16(const FromT* in, Eigen::bfloat16* out,
+                        int num_elements) {
+  std::transform(in, in + num_elements, out, [](FromT a) {
+    return Eigen::bfloat16_impl::float_to_bfloat16_rtne<false>(
+        static_cast<float>(a));
+  });
+}
+
+template <>
+void copyCastToBFloat16(const std::complex<float>* in, Eigen::bfloat16* out,
+                        int num_elements) {
+  std::transform(in, in + num_elements, out, [](std::complex<float> a) {
+    return Eigen::bfloat16_impl::float_to_bfloat16_rtne<false>(std::real(a));
+  });
+}
+
+template <>
+void copyCastToBFloat16(const Eigen::bfloat16* in, Eigen::bfloat16* out,
+                        int num_elements) {
+  std::transform(in, in + num_elements, out,
+                 [](Eigen::bfloat16 a) { return a; });
+}
+
+template <>
+void copyCastToBFloat16(const Eigen::half* in, Eigen::bfloat16* out,
+                        int num_elements) {
+  // half (fp16) -> float -> bfloat16
+  std::transform(in, in + num_elements, out, [](Eigen::half a) {
+    return Eigen::bfloat16_impl::float_to_bfloat16_rtne<false>(
+        Eigen::half_impl::half_to_float(a));
+  });
+}
+
+TfLiteStatus castInt2ToFloat(TfLiteContext* context, const TfLiteTensor* in,
+                             TfLiteTensor* out, int num_elements) {
+  const int8_t* in_data = (const int8_t*)in->data.data;
+  float* out_data = (float*)out->data.data;
+  std::vector<int8_t> unpacked_temp(num_elements);
+  tensor_utils::UnpackPackedIntToInt8(in_data, num_elements, /*bit_width=*/2,
+                                      unpacked_temp.data());
+  for (int i = 0; i < num_elements; ++i) {
+    out_data[i] = static_cast<float>(unpacked_temp[i]);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus castInt4ToFloat(TfLiteContext* context, const TfLiteTensor* in,
+                             TfLiteTensor* out, int num_elements) {
+  const int8_t* in_data = (const int8_t*)in->data.data;
+  float* out_data = (float*)out->data.data;
+  int i = 0;
+#ifdef __ARM_NEON
+  for (; i + 16 <= num_elements / 2; i += 16) {
+    const int8x16_t v0_32 = vld1q_s8(&in_data[i]);
+    const int8x16_t v0_32_low = vshrq_n_s8(vshlq_n_s8(v0_32, 4), 4);
+    const int8x16_t v0_32_high = vshrq_n_s8(v0_32, 4);
+    const int8x16x2_t vzipped = vzipq_s8(v0_32_low, v0_32_high);
+
+    const int16x8_t v0_8 = vmovl_s8(vget_low_s8(vzipped.val[0]));
+    const int16x8_t v8_15 = vmovl_s8(vget_high_s8(vzipped.val[0]));
+    const int16x8_t v16_23 = vmovl_s8(vget_low_s8(vzipped.val[1]));
+    const int16x8_t v24_31 = vmovl_s8(vget_high_s8(vzipped.val[1]));
+
+    const int32x4_t v0_3 = vmovl_s16(vget_low_s16(v0_8));
+    const int32x4_t v4_7 = vmovl_s16(vget_high_s16(v0_8));
+    const int32x4_t v8_11 = vmovl_s16(vget_low_s16(v8_15));
+    const int32x4_t v12_15 = vmovl_s16(vget_high_s16(v8_15));
+    const int32x4_t v16_19 = vmovl_s16(vget_low_s16(v16_23));
+    const int32x4_t v20_23 = vmovl_s16(vget_high_s16(v16_23));
+    const int32x4_t v24_27 = vmovl_s16(vget_low_s16(v24_31));
+    const int32x4_t v28_31 = vmovl_s16(vget_high_s16(v24_31));
+
+    const float32x4_t v0_3_f = vcvtq_f32_s32(v0_3);
+    const float32x4_t v4_7_f = vcvtq_f32_s32(v4_7);
+    const float32x4_t v8_11_f = vcvtq_f32_s32(v8_11);
+    const float32x4_t v12_15_f = vcvtq_f32_s32(v12_15);
+    const float32x4_t v16_19_f = vcvtq_f32_s32(v16_19);
+    const float32x4_t v20_23_f = vcvtq_f32_s32(v20_23);
+    const float32x4_t v24_27_f = vcvtq_f32_s32(v24_27);
+    const float32x4_t v28_31_f = vcvtq_f32_s32(v28_31);
+
+    vst1q_f32(&out_data[i * 2], v0_3_f);
+    vst1q_f32(&out_data[i * 2 + 4], v4_7_f);
+    vst1q_f32(&out_data[i * 2 + 8], v8_11_f);
+    vst1q_f32(&out_data[i * 2 + 12], v12_15_f);
+    vst1q_f32(&out_data[i * 2 + 16], v16_19_f);
+    vst1q_f32(&out_data[i * 2 + 20], v20_23_f);
+    vst1q_f32(&out_data[i * 2 + 24], v24_27_f);
+    vst1q_f32(&out_data[i * 2 + 28], v28_31_f);
+  }
+#endif
+
+  for (; i < (num_elements + 1) / 2; ++i) {
+    int8_t byte = in_data[i];
+    // Shift left first so that sign is properly extended when shifted right
+    int32_t lower = static_cast<int8_t>(byte << 4) >> 4;
+    int32_t higher = byte >> 4;
+    out_data[2 * i] = (float)lower;
+    if (2 * i + 1 < num_elements) {
+      out_data[2 * i + 1] = (float)higher;
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus castUInt4ToFloat(TfLiteContext* context, const TfLiteTensor* in,
+                              TfLiteTensor* out, int num_elements) {
+  const int8_t* in_data = (const int8_t*)in->data.data;
+  float* out_data = (float*)out->data.data;
+  for (int i = 0; i < (num_elements + 1) / 2; ++i) {
+    uint8_t byte = static_cast<uint8_t>(in_data[i]);
+    out_data[2 * i] = static_cast<float>(byte & 0x0F);
+    if (2 * i + 1 < num_elements) {
+      out_data[2 * i + 1] = static_cast<float>(byte >> 4);
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus castFloatToInt4(const float* in, TfLiteTensor* out,
+                             int num_elements) {
+  const float min_val = -8.0f;
+  const float max_val = 7.0f;
+  std::vector<int8_t> unpacked_temp(num_elements);
+  for (int i = 0; i < num_elements; ++i) {
+    unpacked_temp[i] =
+        static_cast<int8_t>(std::max(min_val, std::min(max_val, in[i])));
+  }
+  tensor_utils::PackInt8IntoDenseInt(unpacked_temp.data(), num_elements,
+                                     /*bit_width=*/4, (int8_t*)out->data.data);
+  return kTfLiteOk;
+}
+
+TfLiteStatus castFloatToInt2(const float* in, TfLiteTensor* out,
+                             int num_elements) {
+  const float min_val = -2.0f;
+  const float max_val = 1.0f;
+  std::vector<int8_t> unpacked_temp(num_elements);
+  for (int i = 0; i < num_elements; ++i) {
+    unpacked_temp[i] =
+        static_cast<int8_t>(std::max(min_val, std::min(max_val, in[i])));
+  }
+  tensor_utils::PackInt8IntoDenseInt(unpacked_temp.data(), num_elements,
+                                     /*bit_width=*/2, (int8_t*)out->data.data);
+  return kTfLiteOk;
+}
+
+TfLiteStatus castFloatToUint4(const float* in, TfLiteTensor* out,
+                              int num_elements) {
+  const float min_val = 0.0f;
+  const float max_val = 15.0f;
+  std::vector<int8_t> unpacked_temp(num_elements);
+  for (int i = 0; i < num_elements; ++i) {
+    unpacked_temp[i] =
+        static_cast<int8_t>(std::max(min_val, std::min(max_val, in[i])));
+  }
+  tensor_utils::PackInt8IntoDenseInt(unpacked_temp.data(), num_elements,
+                                     /*bit_width=*/4, (int8_t*)out->data.data);
+  return kTfLiteOk;
+}
+
+template <typename FromT>
+TfLiteStatus copyToTensor(TfLiteContext* context, const FromT* in,
+                          TfLiteTensor* out, int num_elements) {
+  switch (out->type) {
+    case kTfLiteInt64:
+      copyCast(in, out->data.i64, num_elements);
+      break;
+    case kTfLiteInt32:
+      copyCast(in, out->data.i32, num_elements);
+      break;
+    case kTfLiteUInt32:
+      copyCast(in, out->data.u32, num_elements);
+      break;
+    case kTfLiteInt16:
+      copyCast(in, out->data.i16, num_elements);
+      break;
+    case kTfLiteUInt16:
+      copyCast(in, out->data.ui16, num_elements);
+      break;
+    case kTfLiteUInt8:
+      copyCast(in, out->data.uint8, num_elements);
+      break;
+    case kTfLiteInt8:
+      copyCast(in, out->data.int8, num_elements);
+      break;
+    case kTfLiteFloat16:
+      copyCastToFloat16(in, reinterpret_cast<half*>(out->data.f16),
+                        num_elements);
+      break;
+    case kTfLiteBFloat16:
+      copyCastToBFloat16(in, reinterpret_cast<Eigen::bfloat16*>(out->data.bf16),
+                         num_elements);
+      break;
+    case kTfLiteFloat32:
+      copyCast(in, GetTensorData<float>(out), num_elements);
+      break;
+    case kTfLiteFloat64:
+      copyCast(in, out->data.f64, num_elements);
+      break;
+    case kTfLiteBool:
+      copyCast(in, out->data.b, num_elements);
+      break;
+    case kTfLiteComplex64:
+      copyCast(in, reinterpret_cast<std::complex<float>*>(out->data.c64),
+               num_elements);
+      break;
+    case kTfLiteInt4:
+      if (std::is_same<FromT, float>::value) {
+        return castFloatToInt4(reinterpret_cast<const float*>(in), out,
+                               num_elements);
+      } else {
+        TF_LITE_UNSUPPORTED_TYPE(context, out->type, "Cast");
+      }
+    case kTfLiteInt2:
+      if (std::is_same<FromT, float>::value) {
+        return castFloatToInt2(reinterpret_cast<const float*>(in), out,
+                               num_elements);
+      } else {
+        TF_LITE_UNSUPPORTED_TYPE(context, out->type, "Cast");
+      }
+    case kTfLiteUInt4:
+      if (std::is_same<FromT, float>::value) {
+        return castFloatToUint4(reinterpret_cast<const float*>(in), out,
+                                num_elements);
+      } else {
+        TF_LITE_UNSUPPORTED_TYPE(context, out->type, "Cast");
+      }
+    default:
+      // Unsupported type.
+      TF_LITE_UNSUPPORTED_TYPE(context, out->type, "Cast");
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalImpl(TfLiteContext* context, const TfLiteTensor* input,
+                      TfLiteTensor* output, const int num_elements) {
+  switch (input->type) {
+    case kTfLiteInt64:
+      return copyToTensor(context, input->data.i64, output, num_elements);
+    case kTfLiteInt32:
+      return copyToTensor(context, input->data.i32, output, num_elements);
+    case kTfLiteUInt32:
+      return copyToTensor(context, input->data.u32, output, num_elements);
+    case kTfLiteUInt16:
+      return copyToTensor(context, input->data.ui16, output, num_elements);
+    case kTfLiteInt16:
+      return copyToTensor(context, input->data.i16, output, num_elements);
+    case kTfLiteUInt8:
+      return copyToTensor(context, input->data.uint8, output, num_elements);
+    case kTfLiteInt8:
+      return copyToTensor(context, input->data.int8, output, num_elements);
+    case kTfLiteFloat16:
+      return copyToTensor(context,
+                          reinterpret_cast<Eigen::half*>(input->data.f16),
+                          output, num_elements);
+    case kTfLiteBFloat16:
+      return copyToTensor(context,
+                          reinterpret_cast<Eigen::bfloat16*>(input->data.bf16),
+                          output, num_elements);
+    case kTfLiteFloat32:
+      return copyToTensor(context, GetTensorData<float>(input), output,
+                          num_elements);
+    case kTfLiteFloat64:
+      return copyToTensor(context, input->data.f64, output, num_elements);
+    case kTfLiteBool:
+      return copyToTensor(context, input->data.b, output, num_elements);
+    case kTfLiteComplex64:
+      return copyToTensor(
+          context, reinterpret_cast<std::complex<float>*>(input->data.c64),
+          output, num_elements);
+    case kTfLiteInt4:
+      if (output->type != kTfLiteFloat32) {
+        TF_LITE_UNSUPPORTED_TYPE(context, output->type, "Cast");
+      }
+      return castInt4ToFloat(context, input, output, num_elements);
+    case kTfLiteInt2:
+      if (output->type != kTfLiteFloat32) {
+        TF_LITE_UNSUPPORTED_TYPE(context, output->type, "Cast");
+      }
+      return castInt2ToFloat(context, input, output, num_elements);
+    case kTfLiteUInt4:
+      if (output->type != kTfLiteFloat32) {
+        TF_LITE_UNSUPPORTED_TYPE(context, output->type, "Cast");
+      }
+      return castUInt4ToFloat(context, input, output, num_elements);
+    default:
+      // Unsupported type.
+      TF_LITE_UNSUPPORTED_TYPE(context, input->type, "Cast");
+  }
+  return kTfLiteError;
+}
+
+struct OpData {
+  bool cached_output = false;
+};
+
+void* Init(TfLiteContext* context, const char* /*buffer*/, size_t /*length*/) {
+  return new OpData();
+}
+
+void Free(TfLiteContext* context, void* op_data) {
+  delete reinterpret_cast<OpData*>(op_data);
+}
+
+bool OutputCachingEnabled(const TfLiteContext* context) {
+  if (context && context->impl_) {
+    const InterpreterOptions* options =
+        reinterpret_cast<Subgraph*>(context->impl_)->GetOptions();
+    if (options) {
+      return options->GetCacheConstantCastOp();
+    }
+  }
+  return false;
+}
+
+bool ShouldCacheOutput(const TfLiteContext* context,
+                       const TfLiteTensor* input) {
+  return OutputCachingEnabled(context) && IsConstantTensor(input);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // TODO(ahentz): these two checks would make the new implementation
+  // incompatible with some existing models, where params is not specified. It
+  // is OK not to have them because toco would have set input and output types
+  // to match the parameters.
+  // auto* params = reinterpret_cast<TfLiteCastParams*>(node->builtin_data);
+  // TF_LITE_ENSURE_EQ(context, input->type, params->in_data_type);
+  // TF_LITE_ENSURE_EQ(context, output->type, params->out_data_type);
+
+  if (ShouldCacheOutput(context, input)) {
+    output->allocation_type = kTfLiteArenaRwPersistent;
+  }
+
+  TF_LITE_ENSURE_OK(
+      context,
+      context->ResizeTensor(context, output, TfLiteIntArrayCopy(input->dims)));
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  const int num_elements = NumElements(input);
+  TF_LITE_ENSURE_EQ(context, num_elements, NumElements(output));
+
+  OpData& op_data = *reinterpret_cast<OpData*>(node->user_data);
+  if (ShouldCacheOutput(context, input)) {
+    if (op_data.cached_output) {
+      return kTfLiteOk;
+    }
+    op_data.cached_output = true;
+  }
+  return EvalImpl(context, input, output, num_elements);
+}
+
+}  // namespace
+}  // namespace cast
+
+TfLiteRegistration* Register_CAST() {
+  static TfLiteRegistration r = {cast::Init, cast::Free, cast::Prepare,
+                                 cast::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/cast_test.cc

@@ -0,0 +1,564 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include <algorithm>
+#include <complex>
+#include <limits>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/random/random.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/kernels/cast_test_common.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(CastOpModel, CastInt4ToFloat) {
+  CastOpModel m({TensorType_INT4, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.Set4BitInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              Pointwise(FloatingPointEq(), {1.f, 2.f, 3.f, 4.f, 5.f, 6.f}));
+}
+
+TEST(CastOpModel, CastInt4ToFloatLarge) {
+  int num_elements = 40;
+  absl::BitGen bitgen;
+  auto i8rng = [&] {
+    return absl::Uniform<int8_t>(absl::IntervalClosed, bitgen, -8, 7);
+  };
+  std::vector<int8_t> input(num_elements);
+  std::generate(input.begin(), input.end(), i8rng);
+  CastOpModel m({TensorType_INT4, {num_elements}},
+                {TensorType_FLOAT32, {num_elements}});
+  m.Set4BitInput(input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  for (int i = 0; i < input.size(); ++i) {
+    EXPECT_EQ(m.ExtractVector<float>(m.output())[i], input[i]);
+  }
+}
+
+TEST(CastOpModel, CastInt2ToFloat) {
+  CastOpModel m({TensorType_INT2, {2, 4}}, {TensorType_FLOAT32, {2, 4}});
+  m.Set2BitInput({1, 0, -1, -2, 1, 0, -1, -2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              Pointwise(FloatingPointEq(),
+                        {1.f, 0.f, -1.f, -2.f, 1.f, 0.f, -1.f, -2.f}));
+}
+
+TEST(CastOpModel, CastInt2ToFloatLarge) {
+  int num_elements = 40;
+  absl::BitGen bitgen;
+  auto i2rng = [&] {
+    return absl::Uniform<int8_t>(absl::IntervalClosed, bitgen, -2, 1);
+  };
+  std::vector<int8_t> input(num_elements);
+  std::generate(input.begin(), input.end(), i2rng);
+  CastOpModel m({TensorType_INT2, {num_elements}},
+                {TensorType_FLOAT32, {num_elements}});
+  m.Set2BitInput(input);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  for (int i = 0; i < input.size(); ++i) {
+    EXPECT_EQ(m.ExtractVector<float>(m.output())[i], input[i]);
+  }
+}
+
+TEST(CastOpModel, CastFloatToInt4) {
+  CastOpModel m({TensorType_FLOAT32, {2, 4}}, {TensorType_INT4, {2, 4}});
+  m.PopulateTensor<float>(m.input(), {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, -8.f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  TfLiteTensor* output = m.GetOutputTensor(0);
+  int num_elements = NumElements(output);
+  std::vector<int8_t> unpacked_output(num_elements);
+  tensor_utils::UnpackPackedIntToInt8(
+      reinterpret_cast<int8_t*>(output->data.data), num_elements,
+      /*bit_width=*/4, unpacked_output.data());
+  EXPECT_THAT(unpacked_output, ElementsAreArray({1, 2, 3, 4, 5, 6, 7, -8}));
+}
+
+TEST(CastOpModel, CastFloatToInt4Clamp) {
+  CastOpModel m({TensorType_FLOAT32, {1, 4}}, {TensorType_INT4, {1, 4}});
+  m.PopulateTensor<float>(m.input(), {100.f, -100.f, 7.9f, -8.9f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  TfLiteTensor* output = m.GetOutputTensor(0);
+  int num_elements = NumElements(output);
+  std::vector<int8_t> unpacked_output(num_elements);
+  tensor_utils::UnpackPackedIntToInt8(
+      reinterpret_cast<int8_t*>(output->data.data), num_elements,
+      /*bit_width=*/4, unpacked_output.data());
+  EXPECT_THAT(unpacked_output, ElementsAreArray({7, -8, 7, -8}));
+}
+
+TEST(CastOpModel, CastFloatToInt2) {
+  CastOpModel m({TensorType_FLOAT32, {2, 4}}, {TensorType_INT2, {2, 4}});
+  m.PopulateTensor<float>(m.input(),
+                          {1.f, 0.f, -1.f, -2.f, 1.f, 0.f, -1.f, -2.f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  TfLiteTensor* output = m.GetOutputTensor(0);
+  int num_elements = NumElements(output);
+  std::vector<int8_t> unpacked_output(num_elements);
+  tensor_utils::UnpackPackedIntToInt8(
+      reinterpret_cast<int8_t*>(output->data.data), num_elements,
+      /*bit_width=*/2, unpacked_output.data());
+  EXPECT_THAT(unpacked_output, ElementsAreArray({1, 0, -1, -2, 1, 0, -1, -2}));
+}
+
+TEST(CastOpModel, CastFloatToInt2Clamp) {
+  CastOpModel m({TensorType_FLOAT32, {1, 4}}, {TensorType_INT2, {1, 4}});
+  m.PopulateTensor<float>(m.input(), {100.f, -100.f, 1.9f, -2.9f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  TfLiteTensor* output = m.GetOutputTensor(0);
+  int num_elements = NumElements(output);
+  std::vector<int8_t> unpacked_output(num_elements);
+  tensor_utils::UnpackPackedIntToInt8(
+      reinterpret_cast<int8_t*>(output->data.data), num_elements,
+      /*bit_width=*/2, unpacked_output.data());
+  EXPECT_THAT(unpacked_output, ElementsAreArray({1, -2, 1, -2}));
+}
+
+TEST(CastOpModel, CastFloatToUint8Infinity) {
+  CastOpModel m({TensorType_FLOAT32, {2}}, {TensorType_UINT8, {2}});
+  m.PopulateTensor<float>(m.input(), {std::numeric_limits<float>::infinity(),
+                                      -std::numeric_limits<float>::infinity()});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint8_t>(m.output()),
+              ElementsAreArray({std::numeric_limits<uint8_t>::max(),
+                                std::numeric_limits<uint8_t>::min()}));
+}
+
+TEST(CastOpModel, CastFloatToInt16Infinity) {
+  CastOpModel m({TensorType_FLOAT32, {2}}, {TensorType_INT16, {2}});
+  m.PopulateTensor<float>(m.input(), {std::numeric_limits<float>::infinity(),
+                                      -std::numeric_limits<float>::infinity()});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int16_t>(m.output()),
+              ElementsAreArray({std::numeric_limits<int16_t>::max(),
+                                std::numeric_limits<int16_t>::min()}));
+}
+
+TEST(CastOpModel, CastFloatToInt32Infinity) {
+  CastOpModel m({TensorType_FLOAT32, {2}}, {TensorType_INT32, {2}});
+  m.PopulateTensor<float>(m.input(), {std::numeric_limits<float>::infinity(),
+                                      -std::numeric_limits<float>::infinity()});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({std::numeric_limits<int32_t>::max(),
+                                std::numeric_limits<int32_t>::min()}));
+}
+
+TEST(CastOpModel, CastInt16ToFloat) {
+  CastOpModel m({TensorType_INT16, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.PopulateTensor<int16_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<float>(m.output()),
+      Pointwise(FloatingPointEq(), {100.f, 200.f, 300.f, 400.f, 500.f, 600.f}));
+}
+
+TEST(CastOpModel, CastInt16ToInt32) {
+  CastOpModel m({TensorType_INT16, {2, 3}}, {TensorType_INT32, {2, 3}});
+  m.PopulateTensor<int16_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({100, 200, 300, 400, 500, 600}));
+}
+
+TEST(CastOpModel, CastInt32ToFloat) {
+  CastOpModel m({TensorType_INT32, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.PopulateTensor<int32_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<float>(m.output()),
+      Pointwise(FloatingPointEq(), {100.f, 200.f, 300.f, 400.f, 500.f, 600.f}));
+}
+
+TEST(CastOpModel, CastFloatToInt32) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_INT32, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 20.f, 3.f, 0.4f, 0.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({100, 20, 3, 0, 0, 1}));
+}
+
+TEST(CastOpModel, CastFloatToInt16) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_INT16, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 20.f, 3.f, 0.4f, 0.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int16_t>(m.output()),
+              ElementsAreArray({100, 20, 3, 0, 0, 1}));
+}
+
+TEST(CastOpModel, CastInt64ToFloat) {
+  CastOpModel m({TensorType_INT64, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.PopulateTensor<int64_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<float>(m.output()),
+      Pointwise(FloatingPointEq(), {100.f, 200.f, 300.f, 400.f, 500.f, 600.f}));
+}
+
+TEST(CastOpModel, CastFloatToInt64) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_INT64, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 20.f, 3.f, 0.4f, 0.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int64_t>(m.output()),
+              ElementsAreArray({100, 20, 3, 0, 0, 1}));
+}
+
+TEST(CastOpModel, CastFloatToBool) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_BOOL, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, -1.0f, 0.f, 0.4f, 0.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<bool>(m.output()),
+              ElementsAreArray({true, true, false, true, true, true}));
+}
+
+TEST(CastOpModel, CastBoolToFloat) {
+  CastOpModel m({TensorType_BOOL, {3, 2}}, {TensorType_FLOAT32, {3, 2}});
+  m.PopulateTensor<bool>(m.input(), {true, true, false, true, false, true});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              Pointwise(FloatingPointEq(), {1.f, 1.0f, 0.f, 1.0f, 0.0f, 1.0f}));
+}
+
+TEST(CastOpModel, CastFloatToUInt8) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_UINT8, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint8_t>(m.output()),
+              ElementsAreArray({100, 1, 0, 0, 1, 1}));
+}
+
+TEST(CastOpModel, CastUInt8ToFloat) {
+  CastOpModel m({TensorType_UINT8, {3, 2}}, {TensorType_FLOAT32, {3, 2}});
+  m.PopulateTensor<uint8_t>(m.input(), {123, 0, 1, 2, 3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              Pointwise(FloatingPointEq(), {123.f, 0.f, 1.f, 2.f, 3.f, 4.f}));
+}
+
+TEST(CastOpModel, CastFloatToUInt16) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_UINT16, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint16_t>(m.output()),
+              ElementsAreArray({100, 1, 0, 0, 1, 1}));
+}
+
+TEST(CastOpModel, CastUInt16ToFloat) {
+  CastOpModel m({TensorType_UINT16, {3, 2}}, {TensorType_FLOAT32, {3, 2}});
+  m.PopulateTensor<uint16_t>(m.input(), {123, 0, 1, 2, 3, 4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              Pointwise(FloatingPointEq(), {123.f, 0.f, 1.f, 2.f, 3.f, 4.f}));
+}
+
+TEST(CastOpModel, CastInt32ToUInt8) {
+  CastOpModel m({TensorType_INT32, {3, 2}}, {TensorType_UINT8, {3, 2}});
+  m.PopulateTensor<int32_t>(m.input(), {100, 1, 200, 2, 255, 3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint8_t>(m.output()),
+              ElementsAreArray({100, 1, 200, 2, 255, 3}));
+}
+
+TEST(CastOpModel, CastUInt8ToInt32) {
+  CastOpModel m({TensorType_UINT8, {3, 2}}, {TensorType_INT32, {3, 2}});
+  m.PopulateTensor<uint8_t>(m.input(), {100, 1, 200, 2, 255, 3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({100, 1, 200, 2, 255, 3}));
+}
+
+TEST(CastOpModel, CastComplex64ToFloat) {
+  CastOpModel m({TensorType_COMPLEX64, {2, 3}}, {TensorType_FLOAT32, {2, 3}});
+  m.PopulateTensor<std::complex<float>>(
+      m.input(),
+      {std::complex<float>(1.0f, 11.0f), std::complex<float>(2.0f, 12.0f),
+       std::complex<float>(3.0f, 13.0f), std::complex<float>(4.0f, 14.0f),
+       std::complex<float>(5.0f, 15.0f), std::complex<float>(6.0f, 16.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<float>(m.output()),
+      Pointwise(FloatingPointEq(), {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}));
+}
+
+TEST(CastOpModel, CastFloatToComplex64) {
+  CastOpModel m({TensorType_FLOAT32, {2, 3}}, {TensorType_COMPLEX64, {2, 3}});
+  m.PopulateTensor<float>(m.input(), {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<std::complex<float>>(m.output()),
+      ElementsAreArray(
+          {std::complex<float>(1.0f, 0.0f), std::complex<float>(2.0f, 0.0f),
+           std::complex<float>(3.0f, 0.0f), std::complex<float>(4.0f, 0.0f),
+           std::complex<float>(5.0f, 0.0f), std::complex<float>(6.0f, 0.0f)}));
+}
+
+TEST(CastOpModel, CastComplex64ToInt) {
+  CastOpModel m({TensorType_COMPLEX64, {2, 3}}, {TensorType_INT32, {2, 3}});
+  m.PopulateTensor<std::complex<float>>(
+      m.input(),
+      {std::complex<float>(1.0f, 11.0f), std::complex<float>(2.0f, 12.0f),
+       std::complex<float>(3.0f, 13.0f), std::complex<float>(4.0f, 14.0f),
+       std::complex<float>(5.0f, 15.0f), std::complex<float>(6.0f, 16.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int>(m.output()),
+              ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TEST(CastOpModel, CastIntToComplex64) {
+  CastOpModel m({TensorType_INT32, {2, 3}}, {TensorType_COMPLEX64, {2, 3}});
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<std::complex<float>>(m.output()),
+      ElementsAreArray(
+          {std::complex<float>(1.0f, 0.0f), std::complex<float>(2.0f, 0.0f),
+           std::complex<float>(3.0f, 0.0f), std::complex<float>(4.0f, 0.0f),
+           std::complex<float>(5.0f, 0.0f), std::complex<float>(6.0f, 0.0f)}));
+}
+
+TEST(CastOpModel, CastComplex64ToComplex64) {
+  CastOpModel m({TensorType_COMPLEX64, {2, 3}}, {TensorType_COMPLEX64, {2, 3}});
+  m.PopulateTensor<std::complex<float>>(
+      m.input(),
+      {std::complex<float>(1.0f, 11.0f), std::complex<float>(2.0f, 12.0f),
+       std::complex<float>(3.0f, 13.0f), std::complex<float>(4.0f, 14.0f),
+       std::complex<float>(5.0f, 15.0f), std::complex<float>(6.0f, 16.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<std::complex<float>>(m.output()),
+      ElementsAreArray(
+          {std::complex<float>(1.0f, 11.0f), std::complex<float>(2.0f, 12.0f),
+           std::complex<float>(3.0f, 13.0f), std::complex<float>(4.0f, 14.0f),
+           std::complex<float>(5.0f, 15.0f),
+           std::complex<float>(6.0f, 16.0f)}));
+}
+
+TEST(CastOpModel, CastUInt32ToInt32) {
+  CastOpModel m({TensorType_UINT32, {2, 3}}, {TensorType_INT32, {2, 3}});
+  m.PopulateTensor<uint32_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({100, 200, 300, 400, 500, 600}));
+}
+
+TEST(CastOpModel, CastInt32ToUInt32) {
+  CastOpModel m({TensorType_INT32, {2, 3}}, {TensorType_UINT32, {2, 3}});
+  m.PopulateTensor<int32_t>(m.input(), {100, 200, 300, 400, 500, 600});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint32_t>(m.output()),
+              ElementsAreArray({100, 200, 300, 400, 500, 600}));
+}
+
+TEST(CastOpModel, CastUInt8ToInt8) {
+  CastOpModel m({TensorType_UINT8, {2, 3}}, {TensorType_INT8, {2, 3}});
+  m.PopulateTensor<uint8_t>(m.input(), {10, 20, 30, 40, 50, 60});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int8_t>(m.output()),
+              ElementsAreArray({10, 20, 30, 40, 50, 60}));
+}
+
+TEST(CastOpModel, CastInt8ToUInt8) {
+  CastOpModel m({TensorType_INT8, {2, 3}}, {TensorType_UINT8, {2, 3}});
+  m.PopulateTensor<int8_t>(m.input(), {10, 20, 30, 40, 50, 60});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint8_t>(m.output()),
+              ElementsAreArray({10, 20, 30, 40, 50, 60}));
+}
+
+TEST(CastOpModel, CastUInt16ToInt16) {
+  CastOpModel m({TensorType_UINT16, {2, 3}}, {TensorType_INT16, {2, 3}});
+  m.PopulateTensor<uint16_t>(m.input(), {10, 20, 30, 40, 50, 60});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int16_t>(m.output()),
+              ElementsAreArray({10, 20, 30, 40, 50, 60}));
+}
+
+TEST(CastOpModel, CastInt16ToUInt16) {
+  CastOpModel m({TensorType_INT16, {2, 3}}, {TensorType_UINT16, {2, 3}});
+  m.PopulateTensor<int16_t>(m.input(), {10, 20, 30, 40, 50, 60});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<uint16_t>(m.output()),
+              ElementsAreArray({10, 20, 30, 40, 50, 60}));
+}
+
+TEST(CastOpModel, CastFloatToFloat16) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_FLOAT16, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<half>(m.output()),
+      ElementsAreArray({static_cast<half>(100.f), static_cast<half>(1.0f),
+                        static_cast<half>(0.f), static_cast<half>(0.4f),
+                        static_cast<half>(1.999f), static_cast<half>(1.1f)}));
+}
+
+TEST(CastOpModel, CastFloatToBFloat16) {
+  CastOpModel m({TensorType_FLOAT32, {3, 2}}, {TensorType_BFLOAT16, {3, 2}});
+  m.PopulateTensor<float>(m.input(), {100.f, 1.0f, 0.f, 0.4f, 1.999f, 1.1f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<Eigen::bfloat16>(m.output()),
+              ElementsAreArray({static_cast<Eigen::bfloat16>(100.f),
+                                static_cast<Eigen::bfloat16>(1.0f),
+                                static_cast<Eigen::bfloat16>(0.f),
+                                static_cast<Eigen::bfloat16>(0.4f),
+                                static_cast<Eigen::bfloat16>(1.999f),
+                                static_cast<Eigen::bfloat16>(1.1f)}));
+}
+
+TEST(CastOpModel, CastFloat16ToFloat) {
+  CastOpModel m({TensorType_FLOAT16, {3, 2}}, {TensorType_FLOAT32, {3, 2}});
+  m.PopulateTensor<half>(m.input(),
+                         {static_cast<half>(100.f), static_cast<half>(1.0f),
+                          static_cast<half>(0.f), static_cast<half>(0.4f),
+                          static_cast<half>(1.999f), static_cast<half>(1.1f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              ElementsAreArray(ArrayFloatNear(
+                  {100.f, 1.0f, 0.f, 0.399902344f, 1.99902344f, 1.09960938f},
+                  /*max_abs_err=*/0.05f)));
+}
+
+TEST(CastOpModel, CastBFloat16ToFloat) {
+  CastOpModel m({TensorType_BFLOAT16, {3, 2}}, {TensorType_FLOAT32, {3, 2}});
+  m.PopulateTensor<Eigen::bfloat16>(
+      m.input(),
+      {static_cast<Eigen::bfloat16>(100.f), static_cast<Eigen::bfloat16>(1.0f),
+       static_cast<Eigen::bfloat16>(0.f), static_cast<Eigen::bfloat16>(0.4f),
+       static_cast<Eigen::bfloat16>(1.999f),
+       static_cast<Eigen::bfloat16>(1.1)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              ElementsAreArray(ArrayFloatNear(
+                  {100.f, 1.0f, 0.f, 0.400390625f, 2.f, 1.1015625f},
+                  /*max_abs_err=*/0.05f)));
+}
+
+TEST(CastOpModel, CastFloat16ToInt32) {
+  CastOpModel m({TensorType_FLOAT16, {1, 6}}, {TensorType_INT32, {1, 6}});
+  m.PopulateTensor<half>(m.input(),
+                         {static_cast<half>(100.f), static_cast<half>(20.f),
+                          static_cast<half>(3.f), static_cast<half>(0.4f),
+                          static_cast<half>(0.999f), static_cast<half>(1.1f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<int32_t>(m.output()),
+              ElementsAreArray({100, 20, 3, 0, 0, 1}));
+}
+
+TEST(CastOpModel, CastInt32ToFloat16) {
+  CastOpModel m({TensorType_INT32, {1, 6}}, {TensorType_FLOAT16, {1, 6}});
+  m.PopulateTensor<int32_t>(m.input(), {100, 20, 3, 0, 1, -1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.ExtractVector<half>(m.output()),
+      ElementsAreArray({static_cast<half>(100.f), static_cast<half>(20.f),
+                        static_cast<half>(3.f), static_cast<half>(0.f),
+                        static_cast<half>(1.f), static_cast<half>(-1.f)}));
+}
+
+TEST(CastOpModel, CastFloat16ToBFloat16) {
+  CastOpModel m({TensorType_FLOAT16, {1, 6}}, {TensorType_BFLOAT16, {1, 6}});
+  m.PopulateTensor<half>(m.input(),
+                         {static_cast<half>(100.f), static_cast<half>(20.f),
+                          static_cast<half>(3.f), static_cast<half>(0.4f),
+                          static_cast<half>(0.999f), static_cast<half>(1.1f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<Eigen::bfloat16>(m.output()),
+              ElementsAreArray({static_cast<Eigen::bfloat16>(100.f),
+                                static_cast<Eigen::bfloat16>(20.f),
+                                static_cast<Eigen::bfloat16>(3.f),
+                                static_cast<Eigen::bfloat16>(0.4f),
+                                static_cast<Eigen::bfloat16>(0.999f),
+                                static_cast<Eigen::bfloat16>(1.1f)}));
+}
+
+TEST(CastOpModel, CastBFloat16ToFloat16) {
+  CastOpModel m({TensorType_BFLOAT16, {1, 6}}, {TensorType_FLOAT16, {1, 6}});
+  m.PopulateTensor<Eigen::bfloat16>(
+      m.input(),
+      {static_cast<Eigen::bfloat16>(100.f), static_cast<Eigen::bfloat16>(20.f),
+       static_cast<Eigen::bfloat16>(3.f), static_cast<Eigen::bfloat16>(0.4f),
+       static_cast<Eigen::bfloat16>(0.999f),
+       static_cast<Eigen::bfloat16>(1.1f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<half>(m.output()),
+              ElementsAreArray(ArrayFloatNear(
+                  {static_cast<half>(100.f), static_cast<half>(20.f),
+                   static_cast<half>(3.f), static_cast<half>(0.4f),
+                   static_cast<half>(0.999f), static_cast<half>(1.1f)},
+                  /*max_abs_err=*/0.05f)));
+}
+
+TEST(CastOpModel, CastUint4ToFloat) {
+  CastOpModel m({TensorType_UINT4, {1, 6}}, {TensorType_FLOAT32, {1, 6}});
+  m.SetUInt4Input({15, 0, 1, 8, 7, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              ElementsAreArray({15.f, 0.f, 1.f, 8.f, 7.f, 2.f}));
+}
+
+TEST(CastOpModel, CastFloatToUint4) {
+  CastOpModel m({TensorType_FLOAT32, {1, 6}}, {TensorType_UINT4, {1, 6}});
+  m.PopulateTensor<float>(m.input(), {15.f, 0.f, 1.f, 8.f, 7.f, 2.f});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  std::vector<int8_t> unpacked(6);
+  tensor_utils::UnpackPackedIntToInt8(
+      reinterpret_cast<const int8_t*>(m.GetOutputTensor(0)->data.int8), 6,
+      /*bit_width=*/4, unpacked.data(), /*unpack_unsigned=*/true);
+  EXPECT_THAT(unpacked, ElementsAreArray({15, 0, 1, 8, 7, 2}));
+}
+
+TEST(CastOpModel, CastConstInputCachingWorks) {
+  // This tests the implementation of a performance optimization. If that
+  // optimization is changed, this test will likely break/need to be updated.
+  //
+  // We are relying on the fact that casting a constant input can be cached and
+  // that the output tensor does not need to be updated on every call.
+  CastOpModel m({TensorType_INT8, {2, 3}},
+                std::vector<int8_t>{10, 20, 30, 40, 50, 60},
+                {TensorType_FLOAT32, {2, 3}});
+  EXPECT_EQ(m.GetOutputTensor(0)->allocation_type, kTfLiteArenaRwPersistent);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              ElementsAreArray({10, 20, 30, 40, 50, 60}));
+  // We are cheating here. If the values of the output tensor are cached then if
+  // we modify the cache and call the op again the output tensor values should
+  // not change.
+  float* output_data =
+      reinterpret_cast<float*>(m.GetOutputTensor(0)->data.data);
+  for (int i = 0; i < 6; ++i) {
+    ++output_data[i];
+  }
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.ExtractVector<float>(m.output()),
+              ElementsAreArray({11, 21, 31, 41, 51, 61}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/cast_test_common.h

@@ -0,0 +1,78 @@
+/* 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_CAST_TEST_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_CAST_TEST_COMMON_H_
+
+#include <stdint.h>
+
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/interpreter_options.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+using ::testing::ElementsAreArray;
+
+class CastOpModel : public SingleOpModel {
+ public:
+  CastOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_CAST, BuiltinOptions_CastOptions,
+                 CreateCastOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  template <class ConstInputData>
+  CastOpModel(const TensorData& input, ConstInputData&& data,
+              const TensorData& output) {
+    input_ = AddConstInput(input, static_cast<ConstInputData>(data));
+    output_ = AddOutput(output);
+    SetBuiltinOp(BuiltinOperator_CAST, BuiltinOptions_CastOptions,
+                 CreateCastOptions(builder_).Union());
+    BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true, /*allocate_and_delegate=*/false);
+    InterpreterOptions options;
+    options.SetCacheConstantCastOp(true);
+    interpreter_->ApplyOptions(&options);
+    AllocateAndDelegate(/*apply_delegate=*/true);
+  }
+
+  void Set4BitInput(absl::Span<const int8_t> f) {
+    PopulateTensor4bit(input_, 0, f.data(), f.data() + f.size());
+  }
+
+  void SetUInt4Input(absl::Span<const int8_t> f) {
+    PopulateTensor4bit(input_, 0, f.data(), f.data() + f.size());
+  }
+
+  void Set2BitInput(absl::Span<const int8_t> data) {
+    PopulateTensor2bit(input_, 0, data.data(), data.data() + data.size());
+  }
+
+  int input() const { return input_; }
+  int output() const { return output_; }
+
+ protected:
+  int input_;
+  int output_;
+};
+
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_KERNELS_CAST_TEST_COMMON_H_

tensorflow/lite/kernels/ceil.cc

@@ -0,0 +1,70 @@
+/* Copyright 2018 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/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace ceil {
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  output->type = input->type;
+  TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims);
+  return context->ResizeTensor(context, output, output_size);
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  if (input->type != kTfLiteFloat32) {
+    TF_LITE_UNSUPPORTED_TYPE(context, input->type, "Ceil");
+  }
+
+  optimized_ops::Ceil(GetTensorShape(input), GetTensorData<float>(input),
+                      GetTensorShape(output), GetTensorData<float>(output));
+
+  return kTfLiteOk;
+}
+}  // namespace ceil
+
+TfLiteRegistration* Register_CEIL() {
+  static TfLiteRegistration r = {/*init=*/nullptr,
+                                 /*free=*/nullptr, ceil::Prepare, ceil::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/ceil_test.cc

@@ -0,0 +1,79 @@
+/* Copyright 2018 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 <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class CeilOpModel : public SingleOpModel {
+ public:
+  CeilOpModel(std::initializer_list<int> input_shape, TensorType input_type) {
+    input_ = AddInput(TensorType_FLOAT32);
+    output_ = AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_CEIL, BuiltinOptions_NONE, 0);
+    BuildInterpreter({
+        input_shape,
+    });
+  }
+
+  int input() { return input_; }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(CeilOpTest, SingleDim) {
+  CeilOpModel model({2}, TensorType_FLOAT32);
+  model.PopulateTensor<float>(model.input(), {8.5, 0.0});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  EXPECT_THAT(model.GetOutput(), ElementsAreArray({9, 0}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2}));
+}
+
+TEST(CeilOpTest, MultiDims) {
+  CeilOpModel model({2, 1, 1, 5}, TensorType_FLOAT32);
+  std::vector<float> input;
+  if (AllowFp16PrecisionForFp32()) {
+    input = {
+        0.01, 8.01, 0.99, 9.99, 0.5, -0.01, -8.01, -0.99, -9.99, -0.5,
+    };
+  } else {
+    input = {
+        0.0001,  8.0001,  0.9999,  9.9999,  0.5,
+        -0.0001, -8.0001, -0.9999, -9.9999, -0.5,
+    };
+  }
+  model.PopulateTensor<float>(model.input(), input);
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  EXPECT_THAT(model.GetOutput(),
+              Pointwise(FloatingPointEq(), {1, 9, 1, 10, 1, 0, -8, 0, -9, 0}));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2, 1, 1, 5}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/comparisons.cc

@@ -0,0 +1,446 @@
+/* Copyright 2018 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/internal/reference/comparisons.h"
+
+#include <stdint.h>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/string_comparisons.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace comparisons {
+namespace {
+
+constexpr int kInputTensor1 = 0;
+constexpr int kInputTensor2 = 1;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus ComparisonPrepareCommon(TfLiteContext* context, TfLiteNode* node,
+                                     bool is_string_allowed) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  // Don't support string.
+  if (!is_string_allowed) {
+    TF_LITE_ENSURE(context, input1->type != kTfLiteString);
+  }
+  // Currently only support tensors have the same type.
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
+  output->type = kTfLiteBool;
+
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+
+  TfLiteIntArray* output_size = nullptr;
+  if (requires_broadcast) {
+    TF_LITE_ENSURE_OK(context, CalculateShapeForBroadcast(
+                                   context, input1, input2, &output_size));
+  } else {
+    output_size = TfLiteIntArrayCopy(input1->dims);
+  }
+
+  return context->ResizeTensor(context, output, output_size);
+}
+
+TfLiteStatus ComparisonPrepare(TfLiteContext* context, TfLiteNode* node) {
+  return ComparisonPrepareCommon(context, node, false);
+}
+
+TfLiteStatus ComparisonPrepareStringAllowed(TfLiteContext* context,
+                                            TfLiteNode* node) {
+  return ComparisonPrepareCommon(context, node, true);
+}
+
+void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier,
+                        int* left_shift) {
+  if (double_multiplier < 1.0) {
+    QuantizeMultiplierSmallerThanOneExp(double_multiplier, quantized_multiplier,
+                                        left_shift);
+  } else {
+    QuantizeMultiplierGreaterThanOne(double_multiplier, quantized_multiplier,
+                                     left_shift);
+  }
+}
+
+template <typename input_dtype, reference_ops::ComparisonFn<int32> opname>
+void ComparisonQuantized(const TfLiteTensor* input1, const TfLiteTensor* input2,
+                         TfLiteTensor* output, bool requires_broadcast) {
+  if (input1->type == kTfLiteUInt8 || input1->type == kTfLiteInt8 ||
+      input1->type == kTfLiteInt16) {
+    auto input1_offset = -input1->params.zero_point;
+    auto input2_offset = -input2->params.zero_point;
+    const int left_shift = 8;
+
+    int32 input1_multiplier;
+    int32 input2_multiplier;
+    int input1_shift;
+    int input2_shift;
+    QuantizeMultiplier(input1->params.scale, &input1_multiplier, &input1_shift);
+    QuantizeMultiplier(input2->params.scale, &input2_multiplier, &input2_shift);
+
+    ComparisonParams op_params;
+    op_params.left_shift = left_shift;
+    op_params.input1_offset = input1_offset;
+    op_params.input1_multiplier = input1_multiplier;
+    op_params.input1_shift = input1_shift;
+    op_params.input2_offset = input2_offset;
+    op_params.input2_multiplier = input2_multiplier;
+    op_params.input2_shift = input2_shift;
+    if (requires_broadcast) {
+      reference_ops::BroadcastComparison4DSlowWithScaling<input_dtype, opname>(
+          op_params, GetTensorShape(input1), GetTensorData<input_dtype>(input1),
+          GetTensorShape(input2), GetTensorData<input_dtype>(input2),
+          GetTensorShape(output), GetTensorData<bool>(output));
+    } else {
+      reference_ops::ComparisonWithScaling<input_dtype, opname>(
+          op_params, GetTensorShape(input1), GetTensorData<input_dtype>(input1),
+          GetTensorShape(input2), GetTensorData<input_dtype>(input2),
+          GetTensorShape(output), GetTensorData<bool>(output));
+    }
+  }
+}
+
+template <typename T, reference_ops::ComparisonFn<T> opname>
+void Comparison(const TfLiteTensor* input1, const TfLiteTensor* input2,
+                TfLiteTensor* output, bool requires_broadcast) {
+  ComparisonParams op_params;
+  requires_broadcast
+      ? reference_ops::BroadcastComparison4DSlowImpl<T, opname>(
+            op_params, GetTensorShape(input1), GetTensorData<T>(input1),
+            GetTensorShape(input2), GetTensorData<T>(input2),
+            GetTensorShape(output), GetTensorData<bool>(output))
+      : reference_ops::ComparisonImpl<T, opname>(
+            op_params, GetTensorShape(input1), GetTensorData<T>(input1),
+            GetTensorShape(input2), GetTensorData<T>(input2),
+            GetTensorShape(output), GetTensorData<bool>(output));
+}
+
+void ComparisonString(bool (*opname)(const StringRef&, const StringRef&),
+                      const TfLiteTensor* input1, const TfLiteTensor* input2,
+                      TfLiteTensor* output, bool requires_broadcast) {
+  bool* output_data = GetTensorData<bool>(output);
+  if (requires_broadcast) {
+    reference_ops::BroadcastComparison4DSlowStringImpl(
+        opname, GetTensorShape(input1), input1, GetTensorShape(input2), input2,
+        GetTensorShape(output), output_data);
+  } else {
+    reference_ops::ComparisonStringImpl(opname, GetTensorShape(input1), input1,
+                                        GetTensorShape(input2), input2,
+                                        GetTensorShape(output), output_data);
+  }
+}
+
+TfLiteStatus EqualEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  switch (input1->type) {
+    case kTfLiteBool:
+      Comparison<bool, reference_ops::EqualFn>(input1, input2, output,
+                                               requires_broadcast);
+      break;
+    case kTfLiteFloat32:
+      Comparison<float, reference_ops::EqualFn>(input1, input2, output,
+                                                requires_broadcast);
+      break;
+    case kTfLiteFloat16:
+      Comparison<Eigen::half, reference_ops::EqualFn>(input1, input2, output,
+                                                      requires_broadcast);
+      break;
+    case kTfLiteInt16:
+      if (input1->quantization.type == kTfLiteNoQuantization) {
+        Comparison<int16_t, reference_ops::EqualFn>(input1, input2, output,
+                                                    requires_broadcast);
+      } else {
+        ComparisonQuantized<int16_t, reference_ops::EqualFn>(
+            input1, input2, output, requires_broadcast);
+      }
+      break;
+    case kTfLiteInt32:
+      Comparison<int32_t, reference_ops::EqualFn>(input1, input2, output,
+                                                  requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      Comparison<int64_t, reference_ops::EqualFn>(input1, input2, output,
+                                                  requires_broadcast);
+      break;
+    case kTfLiteUInt8:
+      ComparisonQuantized<uint8_t, reference_ops::EqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt8:
+      ComparisonQuantized<int8_t, reference_ops::EqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteString:
+      ComparisonString(reference_ops::StringRefEqualFn, input1, input2, output,
+                       requires_broadcast);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Does not support type %d, requires "
+                         "bool|float|int|uint8|int16|string",
+                         input1->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus NotEqualEval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor1, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputTensor2, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  switch (input1->type) {
+    case kTfLiteBool:
+      Comparison<bool, reference_ops::NotEqualFn>(input1, input2, output,
+                                                  requires_broadcast);
+      break;
+    case kTfLiteFloat32:
+      Comparison<float, reference_ops::NotEqualFn>(input1, input2, output,
+                                                   requires_broadcast);
+      break;
+    case kTfLiteFloat16:
+      Comparison<Eigen::half, reference_ops::NotEqualFn>(input1, input2, output,
+                                                         requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      Comparison<int32_t, reference_ops::NotEqualFn>(input1, input2, output,
+                                                     requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      Comparison<int64_t, reference_ops::NotEqualFn>(input1, input2, output,
+                                                     requires_broadcast);
+      break;
+    case kTfLiteUInt8:
+      ComparisonQuantized<uint8_t, reference_ops::NotEqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt8:
+      ComparisonQuantized<int8_t, reference_ops::NotEqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt16:
+      if (input1->quantization.type != kTfLiteNoQuantization) {
+        ComparisonQuantized<int16_t, reference_ops::NotEqualFn>(
+            input1, input2, output, requires_broadcast);
+      }
+      break;
+    case kTfLiteString:
+      ComparisonString(reference_ops::StringRefNotEqualFn, input1, input2,
+                       output, requires_broadcast);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Does not support type %d, requires "
+                         "bool|float|int|uint8|qint16|string",
+                         input1->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus LessEval(TfLiteContext* context, TfLiteNode* node, int lhs,
+                      int rhs) {
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, lhs, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, rhs, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      Comparison<float, reference_ops::LessFn>(input1, input2, output,
+                                               requires_broadcast);
+      break;
+    case kTfLiteFloat16:
+      Comparison<Eigen::half, reference_ops::LessFn>(input1, input2, output,
+                                                     requires_broadcast);
+      break;
+    case kTfLiteBFloat16:
+      Comparison<Eigen::bfloat16, reference_ops::LessFn>(input1, input2, output,
+                                                         requires_broadcast);
+      break;
+    case kTfLiteInt16:
+      Comparison<int16_t, reference_ops::LessFn>(input1, input2, output,
+                                                 requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      Comparison<int32_t, reference_ops::LessFn>(input1, input2, output,
+                                                 requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      Comparison<int64_t, reference_ops::LessFn>(input1, input2, output,
+                                                 requires_broadcast);
+      break;
+    case kTfLiteUInt8:
+      ComparisonQuantized<uint8_t, reference_ops::LessFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt8:
+      ComparisonQuantized<int8_t, reference_ops::LessFn>(input1, input2, output,
+                                                         requires_broadcast);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Does not support type %d, requires float|int|uint8",
+                         input1->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus LessEqualEval(TfLiteContext* context, TfLiteNode* node, int lhs,
+                           int rhs) {
+  const TfLiteTensor* input1;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, lhs, &input1));
+  const TfLiteTensor* input2;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, rhs, &input2));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  bool requires_broadcast = !HaveSameShapes(input1, input2);
+  switch (input1->type) {
+    case kTfLiteFloat32:
+      Comparison<float, reference_ops::LessEqualFn>(input1, input2, output,
+                                                    requires_broadcast);
+      break;
+    case kTfLiteFloat16:
+      Comparison<Eigen::half, reference_ops::LessEqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt16:
+      Comparison<int16_t, reference_ops::LessEqualFn>(input1, input2, output,
+                                                      requires_broadcast);
+      break;
+    case kTfLiteInt32:
+      Comparison<int32_t, reference_ops::LessEqualFn>(input1, input2, output,
+                                                      requires_broadcast);
+      break;
+    case kTfLiteInt64:
+      Comparison<int64_t, reference_ops::LessEqualFn>(input1, input2, output,
+                                                      requires_broadcast);
+      break;
+    case kTfLiteUInt8:
+      ComparisonQuantized<uint8_t, reference_ops::LessEqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    case kTfLiteInt8:
+      ComparisonQuantized<int8_t, reference_ops::LessEqualFn>(
+          input1, input2, output, requires_broadcast);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context,
+                         "Does not support type %d, requires float|int|uint8",
+                         input1->type);
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus LessEval(TfLiteContext* context, TfLiteNode* node) {
+  return LessEval(context, node, kInputTensor1, kInputTensor2);
+}
+TfLiteStatus LessEqualEval(TfLiteContext* context, TfLiteNode* node) {
+  return LessEqualEval(context, node, kInputTensor1, kInputTensor2);
+}
+
+TfLiteStatus GreaterEval(TfLiteContext* context, TfLiteNode* node) {
+  return LessEval(context, node, kInputTensor2, kInputTensor1);
+}
+TfLiteStatus GreaterEqualEval(TfLiteContext* context, TfLiteNode* node) {
+  return LessEqualEval(context, node, kInputTensor2, kInputTensor1);
+}
+
+}  // namespace
+}  // namespace comparisons
+
+TfLiteRegistration* Register_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepareStringAllowed,
+                                 comparisons::EqualEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_NOT_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepareStringAllowed,
+                                 comparisons::NotEqualEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_GREATER() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::GreaterEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_GREATER_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::GreaterEqualEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_LESS() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, comparisons::ComparisonPrepare, comparisons::LessEval};
+  return &r;
+}
+
+TfLiteRegistration* Register_LESS_EQUAL() {
+  static TfLiteRegistration r = {nullptr, nullptr,
+                                 comparisons::ComparisonPrepare,
+                                 comparisons::LessEqualEval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/comparisons_test.cc

@@ -0,0 +1,900 @@
+/* Copyright 2018 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 <stdint.h>
+
+#include <initializer_list>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAre;
+
+class ComparisonOpModel : public SingleOpModel {
+ public:
+  ComparisonOpModel(std::initializer_list<int> input1_shape,
+                    std::initializer_list<int> input2_shape,
+                    TensorType input_type, BuiltinOperator op,
+                    bool allocate = true) {
+    input1_ = AddInput(input_type);
+    input2_ = AddInput(input_type);
+    output_ = AddOutput(TensorType_BOOL);
+    ConfigureBuiltinOp(op);
+    BuildInterpreter({input1_shape, input2_shape}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true,
+                     /*allocate_and_delegate=*/allocate);
+  }
+
+  ComparisonOpModel(const TensorData& input1, const TensorData& input2,
+                    TensorType input_type, BuiltinOperator op,
+                    bool allocate = true) {
+    input1_ = AddInput(input1);
+    input2_ = AddInput(input2);
+    output_ = AddOutput(TensorType_BOOL);
+    ConfigureBuiltinOp(op);
+    BuildInterpreter({GetShape(input1_), GetShape(input2_)}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true,
+                     /*allocate_and_delegate=*/allocate);
+  }
+
+  int input1() { return input1_; }
+  int input2() { return input2_; }
+
+  std::vector<bool> GetOutput() { return ExtractVector<bool>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input1_;
+  int input2_;
+  int output_;
+
+  void ConfigureBuiltinOp(BuiltinOperator op) {
+    switch (op) {
+      case BuiltinOperator_EQUAL: {
+        SetBuiltinOp(op, BuiltinOptions_EqualOptions,
+                     CreateEqualOptions(builder_).Union());
+        break;
+      }
+      case BuiltinOperator_NOT_EQUAL: {
+        SetBuiltinOp(op, BuiltinOptions_NotEqualOptions,
+                     CreateNotEqualOptions(builder_).Union());
+        break;
+      }
+      case BuiltinOperator_GREATER: {
+        SetBuiltinOp(op, BuiltinOptions_GreaterOptions,
+                     CreateGreaterOptions(builder_).Union());
+        break;
+      }
+      case BuiltinOperator_GREATER_EQUAL: {
+        SetBuiltinOp(op, BuiltinOptions_GreaterEqualOptions,
+                     CreateGreaterEqualOptions(builder_).Union());
+        break;
+      }
+      case BuiltinOperator_LESS: {
+        SetBuiltinOp(op, BuiltinOptions_LessOptions,
+                     CreateLessOptions(builder_).Union());
+        break;
+      }
+      case BuiltinOperator_LESS_EQUAL: {
+        SetBuiltinOp(op, BuiltinOptions_LessEqualOptions,
+                     CreateLessEqualOptions(builder_).Union());
+        break;
+      }
+      default: {
+        FAIL() << "We shouldn't get here.";
+      }
+    }
+  }
+};
+
+template <typename T>
+class FloatComparisonTest : public ::testing::Test {};
+
+using FloatComparisonTestTypes = ::testing::Types<float, half, Eigen::bfloat16>;
+TYPED_TEST_SUITE(FloatComparisonTest, FloatComparisonTestTypes);
+
+TEST(ComparisonsTest, EqualBool) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_BOOL,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<bool>(model.input1(), {true, false, true, false});
+  model.PopulateTensor<bool>(model.input2(), {true, true, false, false});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, Equal) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_EQUAL, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, EqualInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, EqualInt16) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT16,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<int16_t>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int16_t>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, EqualString) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    return;
+  }
+  ComparisonOpModel model({1, 1, 1, 4, 1}, {1, 1, 1, 4, 1}, TensorType_STRING,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<std::string>(model.input1(), {"A", "B", "C", "D"});
+  model.PopulateTensor<std::string>(model.input2(), {"A", "C", "B", "D"});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4, 1));
+}
+
+TEST(ComparisonsTest, EqualBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, EqualBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, false, false, false,
+                                             false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TEST(ComparisonsTest, EqualBroadcastRankFive) {
+  ComparisonOpModel model({1, 2, 1, 2, 2}, {1, 1, 2}, TensorType_INT32,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {1, 2, 3, 4, 5, 6, 7, 8});
+  model.PopulateTensor<int>(model.input2(), {3, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, true, true, false,
+                                             false, false, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 1, 2, 2));
+}
+
+TEST(ComparisonsTest, EqualBroadcastString) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    return;
+  }
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_STRING,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<std::string>(model.input1(), {"A", "B", "A", "B"});
+  model.PopulateTensor<std::string>(model.input2(), {"A"});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, EqualBroadcastStringRankFive) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    return;
+  }
+  ComparisonOpModel model({1, 2, 1, 2, 2}, {1, 1, 2}, TensorType_STRING,
+                          BuiltinOperator_EQUAL);
+  model.PopulateTensor<std::string>(model.input1(),
+                                    {"A", "B", "A", "C", "A", "B", "D", "C"});
+  model.PopulateTensor<std::string>(model.input2(), {"A", "C"});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(true, false, true, true, true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 2, 1, 2, 2));
+}
+
+TEST(ComparisonsTest, NotEqualBool) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_BOOL,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<bool>(model.input1(), {true, false, true, false});
+  model.PopulateTensor<bool>(model.input2(), {true, true, false, false});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, NotEqual) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_NOT_EQUAL, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, NotEqualInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, true, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, NotEqualString) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    return;
+  }
+  ComparisonOpModel model({1, 1, 1, 1, 4}, {1, 1, 1, 1, 4}, TensorType_STRING,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<std::string>(model.input1(), {"A", "B", "C", "D"});
+  model.PopulateTensor<std::string>(model.input2(), {"A", "C", "B", "D"});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, NotEqualBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, true, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, NotEqualBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(true, true, true, true, true, true, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TEST(ComparisonsTest, NotEqualBroadcastString) {
+  if (SingleOpModel::GetForceUseNnapi()) {
+    return;
+  }
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_STRING,
+                          BuiltinOperator_NOT_EQUAL);
+  model.PopulateTensor<std::string>(model.input1(), {"A", "B", "A", "B"});
+  model.PopulateTensor<std::string>(model.input2(), {"A"});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, Greater) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_GREATER, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_GREATER);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, false, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_GREATER);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, false, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_GREATER);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(false, true, true, false, false, true, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, GreaterEqual) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_GREATER_EQUAL, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterEqualInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_GREATER_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterEqualInt16) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT16,
+                          BuiltinOperator_GREATER_EQUAL);
+  model.PopulateTensor<int16_t>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int16_t>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterEqualBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_GREATER_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, GreaterEqualBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_GREATER_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(false, true, true, false, false, true, true, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, Less) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_LESS, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_LESS);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessInt16) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT16,
+                          BuiltinOperator_LESS);
+  model.PopulateTensor<int16_t>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int16_t>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_LESS);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_LESS);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 6, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(true, false, false, true, true, false, false, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TYPED_TEST(FloatComparisonTest, LessEqual) {
+  using T = TypeParam;
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, GetTensorType<T>(),
+                          BuiltinOperator_LESS_EQUAL, /*allocate=*/false);
+  TFLITE_ALLOCATE_AND_CHECK(T, &model);
+  model.template PopulateTensor<T>(model.input1(), {0.1, 0.9, 0.7, 0.3});
+  model.template PopulateTensor<T>(model.input2(), {0.1, 0.2, 0.6, 0.5});
+  TFLITE_INVOKE_AND_CHECK(T, &model);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessEqualInt) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_LESS_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, true, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessEqualBroadcast) {
+  ComparisonOpModel model({1, 1, 1, 4}, {1, 1, 1, 1}, TensorType_INT32,
+                          BuiltinOperator_LESS_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3});
+  model.PopulateTensor<int>(model.input2(), {7});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, true, true));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 1, 4));
+}
+
+TEST(ComparisonsTest, LessEqualBroadcastTwoD) {
+  ComparisonOpModel model({1, 1, 2, 4}, {1, 1, 1, 4}, TensorType_INT32,
+                          BuiltinOperator_LESS_EQUAL);
+  model.PopulateTensor<int>(model.input1(), {-1, 9, 7, 3, 2, 4, 2, 8});
+  model.PopulateTensor<int>(model.input2(), {7, 1, 2, 4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(),
+              ElementsAre(true, false, false, true, true, false, true, false));
+  EXPECT_THAT(model.GetOutputShape(), ElementsAre(1, 1, 2, 4));
+}
+
+TEST(QuantizedComparisonsTest, EqualUInt8Quantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_EQUAL);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, true, false));
+}
+
+TEST(QuantizedComparisonsTest, EqualInt8Quantized) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  ComparisonOpModel model({TensorType_INT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_INT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_INT8, BuiltinOperator_EQUAL);
+  model.QuantizeAndPopulate<int8_t>(model.input1(), {1, -9, 7, 3});
+  model.QuantizeAndPopulate<int8_t>(model.input2(), {-1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, true, false));
+}
+
+TEST(QuantizedComparisonsTest, EqualInt16Quantized) {
+  const float kMin = std::numeric_limits<int16_t>::min() + 1;
+  const float kMax = std::numeric_limits<int16_t>::max();
+  ComparisonOpModel model({TensorType_INT16, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_INT16, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_INT16, BuiltinOperator_EQUAL);
+  model.QuantizeAndPopulate<int16_t>(model.input1(), {10, -90, 70, kMin});
+  model.QuantizeAndPopulate<int16_t>(model.input2(), {10, 20, 71, kMin});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+}
+
+TEST(QuantizedComparisonsTest, NotEqualUInt8Quantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_NOT_EQUAL);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 7, 0});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, false, true));
+}
+
+TEST(QuantizedComparisonsTest, NotEqualInt8Quantized) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  ComparisonOpModel model({TensorType_INT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_INT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_INT8, BuiltinOperator_NOT_EQUAL);
+  model.QuantizeAndPopulate<int8_t>(model.input1(), {1, -9, 7, 3});
+  model.QuantizeAndPopulate<int8_t>(model.input2(), {1, 2, 7, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, false, true));
+}
+
+TEST(QuantizedComparisonsTest, NotEqualInt16Quantized) {
+  const float kMin = std::numeric_limits<int16_t>::min() + 1;
+  const float kMax = std::numeric_limits<int16_t>::max();
+  ComparisonOpModel model({TensorType_INT16, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_INT16, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_INT16, BuiltinOperator_NOT_EQUAL);
+  model.QuantizeAndPopulate<int16_t>(model.input1(), {10, -90, 70, kMin + 1});
+  model.QuantizeAndPopulate<int16_t>(model.input2(), {10, 20, 71, kMin + 2});
+  model.Invoke();
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, true));
+}
+
+TEST(ComparisonsTest, GreaterQuantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_GREATER);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+}
+
+TEST(ComparisonsTest, GreaterQuantizedSmallRange) {
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, 0.0, 1.0},
+                          {TensorType_UINT8, {1, 2, 2, 1}, 0.0, 2.0},
+                          TensorType_UINT8, BuiltinOperator_GREATER);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1.0, 0.5, 0.35, 0.1});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1.01, 0.25, 0.3, 0.4});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, true, true, false));
+}
+
+TEST(ComparisonsTest, GreaterEqualQuantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_GREATER_EQUAL);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, true, true, false));
+}
+
+TEST(ComparisonsTest, LessQuantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_LESS);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(false, false, false, true));
+}
+
+TEST(ComparisonsTest, LessEqualQuantized) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  ComparisonOpModel model({TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          {TensorType_UINT8, {1, 2, 2, 1}, kMin, kMax},
+                          TensorType_UINT8, BuiltinOperator_LESS_EQUAL);
+  model.QuantizeAndPopulate<uint8_t>(model.input1(), {1, 9, 7, 3});
+  model.QuantizeAndPopulate<uint8_t>(model.input2(), {1, 2, 6, 5});
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(model.GetOutput(), ElementsAre(true, false, false, true));
+}
+
+TEST(ComparisonsTest, QuantizedEqualWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_EQUAL);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {2});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(false, true, false, false, false, false))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedUInt8NotEqualWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_NOT_EQUAL);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {2});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, true, true, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedInt8NotEqualWithBroadcast) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_NOT_EQUAL);
+    model.QuantizeAndPopulate<int8_t>(model.input1(), {-20, 2, 7, -8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {2});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, true, true, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedUInt8GreaterWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_GREATER);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, false, false, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedInt8GreaterWithBroadcast) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_GREATER);
+    model.QuantizeAndPopulate<int8_t>(model.input1(), {20, -2, -71, 8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, false, false, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest,
+     QuantizedInt8GreaterWithBroadcastMultiplierGreaterThanOne) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_GREATER);
+    model.QuantizeAndPopulate<int8_t>(model.input1(),
+                                      {572, -2, -71, 8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, false, false, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedUInt8GreaterEqualWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_GREATER_EQUAL);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, false, true, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedInt8GreaterEqualWithBroadcast) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_GREATER_EQUAL);
+    model.QuantizeAndPopulate<int8_t>(model.input1(), {20, -2, -71, 8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(true, false, false, true, true, true))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedUInt8LessWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_LESS);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(false, true, true, false, false, false))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedInt8LessWithBroadcast) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_LESS);
+    model.QuantizeAndPopulate<int8_t>(model.input1(), {20, -2, -71, 8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(false, true, true, false, false, false))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedUInt8LessEqualWithBroadcast) {
+  const float kMin = -1.f;
+  const float kMax = 128.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_UINT8, test_shapes[i], kMin, kMax},
+                            {TensorType_UINT8, {}, kMin, kMax},
+                            TensorType_UINT8, BuiltinOperator_LESS_EQUAL);
+    model.QuantizeAndPopulate<uint8_t>(model.input1(), {20, 2, 7, 8, 11, 20});
+    model.QuantizeAndPopulate<uint8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(false, true, true, true, false, false))
+        << "With shape number " << i;
+  }
+}
+
+TEST(ComparisonsTest, QuantizedInt8LessEqualWithBroadcast) {
+  const float kMin = -127.f;
+  const float kMax = 127.f;
+  std::vector<std::vector<int>> test_shapes = {
+      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
+  for (int i = 0; i < test_shapes.size(); ++i) {
+    ComparisonOpModel model({TensorType_INT8, test_shapes[i], kMin, kMax},
+                            {TensorType_INT8, {}, kMin, kMax}, TensorType_INT8,
+                            BuiltinOperator_LESS_EQUAL);
+    model.QuantizeAndPopulate<int8_t>(model.input1(), {20, -2, -71, 8, 11, 20});
+    model.QuantizeAndPopulate<int8_t>(model.input2(), {8});
+    ASSERT_EQ(model.Invoke(), kTfLiteOk);
+    EXPECT_THAT(model.GetOutput(),
+                ElementsAre(false, true, true, true, false, false))
+        << "With shape number " << i;
+  }
+}
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/complex_support.cc

@@ -0,0 +1,180 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <complex>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace complex {
+
+static const int kInputTensor = 0;
+static const int kOutputTensor = 0;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+
+  TF_LITE_ENSURE(context, input->type == kTfLiteComplex64 ||
+                              input->type == kTfLiteComplex128);
+
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  if (input->type == kTfLiteComplex64) {
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat32);
+  } else {
+    TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteFloat64);
+  }
+
+  TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
+  return context->ResizeTensor(context, output, output_shape);
+}
+
+template <typename T, typename ExtractF>
+void ExtractData(const TfLiteTensor* input, ExtractF extract_func,
+                 TfLiteTensor* output) {
+  const std::complex<T>* input_data = GetTensorData<std::complex<T>>(input);
+  T* output_data = GetTensorData<T>(output);
+  const int input_size = NumElements(input);
+  for (int i = 0; i < input_size; ++i) {
+    *output_data++ = extract_func(*input_data++);
+  }
+}
+
+TfLiteStatus EvalReal(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  switch (input->type) {
+    case kTfLiteComplex64: {
+      ExtractData<float>(
+          input,
+          static_cast<float (*)(const std::complex<float>&)>(std::real<float>),
+          output);
+      break;
+    }
+    case kTfLiteComplex128: {
+      ExtractData<double>(input,
+                          static_cast<double (*)(const std::complex<double>&)>(
+                              std::real<double>),
+                          output);
+      break;
+    }
+    default: {
+      TF_LITE_KERNEL_LOG(context,
+                         "Unsupported input type, Real op only supports "
+                         "complex input, but got: %s",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalImag(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  switch (input->type) {
+    case kTfLiteComplex64: {
+      ExtractData<float>(
+          input,
+          static_cast<float (*)(const std::complex<float>&)>(std::imag<float>),
+          output);
+      break;
+    }
+    case kTfLiteComplex128: {
+      ExtractData<double>(input,
+                          static_cast<double (*)(const std::complex<double>&)>(
+                              std::imag<double>),
+                          output);
+      break;
+    }
+    default: {
+      TF_LITE_KERNEL_LOG(context,
+                         "Unsupported input type, Imag op only supports "
+                         "complex input, but got: %s",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalAbs(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  switch (input->type) {
+    case kTfLiteComplex64: {
+      ExtractData<float>(
+          input,
+          static_cast<float (*)(const std::complex<float>&)>(std::abs<float>),
+          output);
+      break;
+    }
+    case kTfLiteComplex128: {
+      ExtractData<double>(input,
+                          static_cast<double (*)(const std::complex<double>&)>(
+                              std::abs<double>),
+                          output);
+      break;
+    }
+    default: {
+      TF_LITE_KERNEL_LOG(context,
+                         "Unsupported input type, ComplexAbs op only supports "
+                         "complex input, but got: %s",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace complex
+
+TfLiteRegistration* Register_REAL() {
+  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+                                 complex::Prepare, complex::EvalReal};
+  return &r;
+}
+
+TfLiteRegistration* Register_IMAG() {
+  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+                                 complex::Prepare, complex::EvalImag};
+  return &r;
+}
+
+TfLiteRegistration* Register_COMPLEX_ABS() {
+  static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr,
+                                 complex::Prepare, complex::EvalAbs};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/complex_support_test.cc

@@ -0,0 +1,233 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <complex>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace {
+
+template <typename T>
+class RealOpModel : public SingleOpModel {
+ public:
+  RealOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+
+    output_ = AddOutput(output);
+
+    const std::vector<uint8_t> custom_option;
+    SetBuiltinOp(BuiltinOperator_REAL, BuiltinOptions_NONE, 0);
+
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input() { return input_; }
+
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(RealOpTest, SimpleFloatTest) {
+  RealOpModel<float> m({TensorType_COMPLEX64, {2, 4}},
+                       {TensorType_FLOAT32, {}});
+
+  m.PopulateTensor<std::complex<float>>(m.input(), {{75, 0},
+                                                    {-6, -1},
+                                                    {9, 0},
+                                                    {-10, 5},
+                                                    {-3, 2},
+                                                    {-6, 11},
+                                                    {0, 0},
+                                                    {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {75, -6, 9, -10, -3, -6, 0, 22.1f})));
+}
+
+TEST(RealOpTest, SimpleDoubleTest) {
+  RealOpModel<double> m({TensorType_COMPLEX128, {2, 4}},
+                        {TensorType_FLOAT64, {}});
+
+  m.PopulateTensor<std::complex<double>>(m.input(), {{75, 0},
+                                                     {-6, -1},
+                                                     {9, 0},
+                                                     {-10, 5},
+                                                     {-3, 2},
+                                                     {-6, 11},
+                                                     {0, 0},
+                                                     {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {75, -6, 9, -10, -3, -6, 0, 22.1f})));
+}
+
+template <typename T>
+class ImagOpModel : public SingleOpModel {
+ public:
+  ImagOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+
+    output_ = AddOutput(output);
+
+    const std::vector<uint8_t> custom_option;
+    SetBuiltinOp(BuiltinOperator_IMAG, BuiltinOptions_NONE, 0);
+
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input() { return input_; }
+
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(ImagOpTest, SimpleFloatTest) {
+  ImagOpModel<float> m({TensorType_COMPLEX64, {2, 4}},
+                       {TensorType_FLOAT32, {}});
+
+  m.PopulateTensor<std::complex<float>>(m.input(), {{75, 7},
+                                                    {-6, -1},
+                                                    {9, 3.5},
+                                                    {-10, 5},
+                                                    {-3, 2},
+                                                    {-6, 11},
+                                                    {0, 0},
+                                                    {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {7, -1, 3.5f, 5, 2, 11, 0, 33.3f})));
+}
+
+TEST(ImagOpTest, SimpleDoubleTest) {
+  ImagOpModel<double> m({TensorType_COMPLEX128, {2, 4}},
+                        {TensorType_FLOAT64, {}});
+
+  m.PopulateTensor<std::complex<double>>(m.input(), {{75, 7},
+                                                     {-6, -1},
+                                                     {9, 3.5},
+                                                     {-10, 5},
+                                                     {-3, 2},
+                                                     {-6, 11},
+                                                     {0, 0},
+                                                     {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {7, -1, 3.5f, 5, 2, 11, 0, 33.3f})));
+}
+
+template <typename T>
+class ComplexAbsOpModel : public SingleOpModel {
+ public:
+  ComplexAbsOpModel(const TensorData& input, const TensorData& output) {
+    input_ = AddInput(input);
+
+    output_ = AddOutput(output);
+
+    const std::vector<uint8_t> custom_option;
+    SetBuiltinOp(BuiltinOperator_COMPLEX_ABS, BuiltinOptions_NONE, 0);
+
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  int input() { return input_; }
+
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(ComplexAbsOpTest, IncompatibleType64Test) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      ComplexAbsOpModel<float> m({TensorType_COMPLEX64, {2, 4}},
+                                 {TensorType_FLOAT64, {}}),
+      "output->type != kTfLiteFloat32");
+}
+
+TEST(ComplexAbsOpTest, IncompatibleType128Test) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      ComplexAbsOpModel<float> m({TensorType_COMPLEX128, {2, 4}},
+                                 {TensorType_FLOAT32, {}}),
+      "output->type != kTfLiteFloat64");
+}
+
+TEST(ComplexAbsOpTest, SimpleFloatTest) {
+  ComplexAbsOpModel<float> m({TensorType_COMPLEX64, {2, 4}},
+                             {TensorType_FLOAT32, {}});
+
+  m.PopulateTensor<std::complex<float>>(m.input(), {{75, 7},
+                                                    {-6, -1},
+                                                    {9, 3.5},
+                                                    {-10, 5},
+                                                    {-3, 2},
+                                                    {-6, 11},
+                                                    {0, 0},
+                                                    {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), testing::ElementsAre(2, 4));
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {75.32596f, 6.0827627f, 9.656604f, 11.18034f,
+                                  3.6055512f, 12.529964f, 0.f, 39.966236f})));
+}
+
+TEST(ComplexAbsOpTest, SimpleDoubleTest) {
+  ComplexAbsOpModel<double> m({TensorType_COMPLEX128, {2, 4}},
+                              {TensorType_FLOAT64, {}});
+
+  m.PopulateTensor<std::complex<double>>(m.input(), {{75, 7},
+                                                     {-6, -1},
+                                                     {9, 3.5},
+                                                     {-10, 5},
+                                                     {-3, 2},
+                                                     {-6, 11},
+                                                     {0, 0},
+                                                     {22.1, 33.3}});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), testing::ElementsAre(2, 4));
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(ArrayFloatNear(
+                                 {75.32596f, 6.0827627f, 9.656604f, 11.18034f,
+                                  3.6055512f, 12.529964f, 0.f, 39.966236f})));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/concatenation.cc

@@ -0,0 +1,295 @@
+/* Copyright 2017 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/internal/reference/concatenation.h"
+
+#include <stdint.h>
+
+#include <cstddef>
+#include <cstring>
+#include <limits>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace concatenation {
+
+// This file has two implementation of Concatenation.
+enum KernelType {
+  kReference,
+  kGenericOptimized,
+};
+
+template <KernelType kernel_type>
+TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node, int axis,
+                      TfLiteTensor* output) {
+// TODO(ahentz): Creating 'all_inputs' below is not very efficient. We should
+// allocate and populate these during Prepare().
+// TODO(ycling): Activation function parameter is ignored. For now we don't have
+// a model with a Concatenation with fused activation function.
+#define TF_LITE_CONCATENATION(scalar)                                         \
+  {                                                                           \
+    VectorOfTensors<scalar> all_inputs(*context, *node->inputs);              \
+    tflite::ConcatenationParams op_params;                                    \
+    op_params.axis = axis;                                                    \
+    op_params.inputs_count = node->inputs->size;                              \
+    if (kernel_type == kReference) {                                          \
+      reference_ops::Concatenation(op_params, all_inputs.shapes(),            \
+                                   all_inputs.data(), GetTensorShape(output), \
+                                   GetTensorData<scalar>(output));            \
+    } else {                                                                  \
+      optimized_ops::Concatenation(op_params, all_inputs.shapes(),            \
+                                   all_inputs.data(), GetTensorShape(output), \
+                                   GetTensorData<scalar>(output));            \
+    }                                                                         \
+  }
+
+#define TF_LITE_CONCATENATION_QUANTIZED()                         \
+  {                                                               \
+    VectorOfQuantizedTensors all_inputs(*context, *node->inputs); \
+    tflite::ConcatenationParams op_params;                        \
+    op_params.axis = axis;                                        \
+    op_params.input_zeropoint = all_inputs.zero_point();          \
+    op_params.input_scale = all_inputs.scale();                   \
+    op_params.inputs_count = node->inputs->size;                  \
+    op_params.output_zeropoint = output->params.zero_point;       \
+    op_params.output_scale = output->params.scale;                \
+    if (kernel_type == kReference) {                              \
+      reference_ops::ConcatenationWithScaling(                    \
+          op_params, all_inputs.shapes(), all_inputs.data(),      \
+          GetTensorShape(output), GetTensorData<uint8>(output));  \
+    } else {                                                      \
+      optimized_ops::ConcatenationWithScaling(                    \
+          op_params, all_inputs.shapes(), all_inputs.data(),      \
+          GetTensorShape(output), GetTensorData<uint8>(output));  \
+    }                                                             \
+  }
+
+  switch (output->type) {
+    case kTfLiteUInt8:
+      TF_LITE_CONCATENATION_QUANTIZED();
+      return kTfLiteOk;
+    case kTfLiteInt4:
+      TF_LITE_CONCATENATION(Int4);
+      return kTfLiteOk;
+    default:
+      break;
+  }
+
+  switch (TfLiteTypeGetSizeBits(output->type)) {
+    case 8:
+      TF_LITE_CONCATENATION(int8_t);
+      break;
+    case 16:
+      TF_LITE_CONCATENATION(int16_t);
+      break;
+    case 32:
+      TF_LITE_CONCATENATION(int32_t);
+      break;
+    case 64:
+      TF_LITE_CONCATENATION(int64_t);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type '%s' is not supported currently.",
+                         TfLiteTypeGetName(output->type));
+      return kTfLiteError;
+  }
+
+#undef TF_LITE_CONCATENATION_QUANTIZED
+#undef TF_LITE_CONCATENATION
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteConcatenationParams*>(node->builtin_data);
+  int axis = params->axis;
+  int num_inputs = node->inputs->size;
+
+  // The number of dimensions of the input tensors must match, and all
+  // dimensions except 'axis' must be equal.
+  const TfLiteTensor* t0;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &t0));
+  TfLiteType input_type = t0->type;
+  if (axis < 0) axis += t0->dims->size;
+  TF_LITE_ENSURE(context, axis >= 0);
+  TF_LITE_ENSURE(context,
+                 axis < t0->dims->size || (t0->dims->size == 0 && axis == 0));
+
+  TF_LITE_ENSURE_EQ(context, params->activation, kTfLiteActNone);
+  TF_LITE_ENSURE(context,
+                 input_type == kTfLiteFloat32 || input_type == kTfLiteFloat16 ||
+                     input_type == kTfLiteBFloat16 ||
+                     input_type == kTfLiteUInt8 || input_type == kTfLiteInt8 ||
+                     input_type == kTfLiteInt16 || input_type == kTfLiteInt32 ||
+                     input_type == kTfLiteInt64 || input_type == kTfLiteBool ||
+                     input_type == kTfLiteUInt32 || input_type == kTfLiteInt4);
+
+  // Check to see if we can calculate the output now.
+  bool all_inputs_at_prepare = true;
+  for (int i = 0; i < num_inputs; ++i) {
+    const TfLiteTensor* t;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &t));
+    if (!IsConstantOrPersistentTensor(t)) {
+      all_inputs_at_prepare = false;
+      break;
+    }
+  }
+  // Output dimensions will match input dimensions, except 'axis', which
+  // will be the sum of inputs
+  int sum_axis = t0->dims->size > 0 ? t0->dims->data[axis] : 1;
+  // Check if we are concatenating constant scalars.
+  if (all_inputs_at_prepare && t0->dims->size == 0 && axis == 0 &&
+      input_type != kTfLiteInt4) {
+    for (int i = 1; i < num_inputs; ++i) {
+      const TfLiteTensor* t;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &t));
+      TF_LITE_ENSURE_EQ(context, t->dims->size, t0->dims->size);
+    }
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+    TfLiteIntArray* output_size = TfLiteIntArrayCreate(1);
+    output_size->data[0] = num_inputs;
+    SetTensorToPersistentRo(output);
+    context->ResizeTensor(context, output, output_size);
+    size_t input_type_size;
+    TF_LITE_ENSURE_STATUS(GetSizeOfType(context, t0->type, &input_type_size));
+    void* o_data = output->data.data;
+    for (int i = 0; i < num_inputs; ++i) {
+      const TfLiteTensor* t;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &t));
+      const void* i_data = t->data.data;
+      memcpy(o_data, i_data, input_type_size);
+      o_data = (void*)((uintptr_t)o_data + input_type_size);
+    }
+    return kTfLiteOk;
+  } else {
+    // Scalar concatenation is only supported for constant tensors
+    // (handled above).
+    TF_LITE_ENSURE(context, t0->dims->size > 0);
+    for (int i = 1; i < num_inputs; ++i) {
+      const TfLiteTensor* t;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &t));
+      TF_LITE_ENSURE_EQ(context, t->dims->size, t0->dims->size);
+      TF_LITE_ENSURE_EQ(context, t->type, input_type);
+      for (int d = 0; d < t0->dims->size; ++d) {
+        if (d == axis) {
+          // Avoid integer overflow in sum_axis below
+          TF_LITE_ENSURE(context, t->dims->data[axis] >= 0);
+          TF_LITE_ENSURE(context,
+                         t->dims->data[axis] <=
+                             std::numeric_limits<int>::max() - sum_axis);
+          sum_axis += t->dims->data[axis];
+        } else {
+          TF_LITE_ENSURE_EQ(context, t->dims->data[d], t0->dims->data[d]);
+        }
+      }
+    }
+  }
+
+  TfLiteIntArray* output_size = TfLiteIntArrayCreate(t0->dims->size);
+  for (int d = 0; d < t0->dims->size; ++d) {
+    output_size->data[d] = (d == axis) ? sum_axis : t0->dims->data[d];
+  }
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input_type);
+
+  if (input_type == kTfLiteInt8) {
+    // Make sure there is no re-scaling needed for Int8 quantized kernel. This
+    // is a restriction we introduced to Int8 kernels.
+    VectorOfTensors<int8_t> all_inputs(*context, *node->inputs);
+    for (int i = 0; i < node->inputs->size; ++i) {
+      const TfLiteTensor* t;
+      TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &t));
+      TF_LITE_ENSURE_EQ(context, t->params.scale, output->params.scale);
+      TF_LITE_ENSURE_EQ(context, t->params.zero_point,
+                        output->params.zero_point);
+    }
+  }
+
+  if (input_type == kTfLiteInt16) {
+    // Make sure that all Int16 inputs have a null zero-point.
+    for (int i = 0; i < node->inputs->size; ++i) {
+      const TfLiteTensor* t = GetInput(context, node, i);
+      TF_LITE_ENSURE_EQ(context, t->params.zero_point, 0);
+    }
+    TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
+  }
+
+  if (all_inputs_at_prepare) {
+    SetTensorToPersistentRo(output);
+    context->ResizeTensor(context, output, output_size);
+    return EvalImpl<kReference>(context, node, axis, output);
+  }
+  return context->ResizeTensor(context, output, output_size);
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteConcatenationParams*>(node->builtin_data);
+  int axis = params->axis;
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  if (IsConstantOrPersistentTensor(output)) {
+    // Output is computed in Prepare.
+    return kTfLiteOk;
+  }
+  if (axis < 0) axis += output->dims->size;
+
+  return EvalImpl<kernel_type>(context, node, axis, output);
+}
+
+#undef TF_LITE_MACRO_DISPATCH
+
+}  // namespace concatenation
+
+TfLiteRegistration* Register_CONCATENATION_REF() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, concatenation::Prepare,
+      concatenation::Eval<concatenation::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONCATENATION_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, concatenation::Prepare,
+      concatenation::Eval<concatenation::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONCATENATION() {
+  // TODO(ahentz): It turns out the two versions of Concatenation are almost
+  // identical, so we should consider removing one.
+  return Register_CONCATENATION_GENERIC_OPT();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/concatenation_test.cc

@@ -0,0 +1,935 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include <initializer_list>
+#include <limits>
+#include <type_traits>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/types/half.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class BaseConcatenationOpModel : public SingleOpModel {
+ public:
+  // TODO(ahentz): Also test different activation types, axis, input
+  // dimensions.
+  BaseConcatenationOpModel() {}
+  BaseConcatenationOpModel(const std::vector<TensorData>& input_template,
+                           int axis, int num_inputs,
+                           const TensorData& output_template) {
+    std::vector<std::vector<int>> all_input_shapes;
+    CHECK_EQ(input_template.size(), num_inputs);
+    for (int i = 0; i < num_inputs; ++i) {
+      all_input_shapes.push_back(input_template[i].shape);
+      AddInput(input_template[i]);
+    }
+    output_ = AddOutput({output_template.type, /*shape=*/{},
+                         output_template.min, output_template.max});
+    SetBuiltinOp(
+        BuiltinOperator_CONCATENATION, BuiltinOptions_ConcatenationOptions,
+        CreateConcatenationOptions(builder_, axis, ActivationFunctionType_NONE)
+            .Union());
+    BuildInterpreter(all_input_shapes);
+  }
+  BaseConcatenationOpModel(const TensorData& input_template, int axis,
+                           int num_inputs)
+      : BaseConcatenationOpModel(
+            std::vector<TensorData>(num_inputs, input_template), axis,
+            num_inputs, input_template) {}
+
+ protected:
+  int output_;
+};
+
+template <typename T>
+class ConcatenationOpModel : public BaseConcatenationOpModel {
+ public:
+  using BaseConcatenationOpModel::BaseConcatenationOpModel;
+  void SetInput(int index, std::initializer_list<T> data) {
+    PopulateTensor<T>(index, data);
+  }
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+};
+
+class QuantizedConcatenationOpModel : public BaseConcatenationOpModel {
+ public:
+  using BaseConcatenationOpModel::BaseConcatenationOpModel;
+
+  template <typename T>
+  void SetInput(int index, std::initializer_list<float> data) {
+    QuantizeAndPopulate<T>(index, data);
+  }
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+  template <typename T>
+  std::vector<float> GetDequantizedOutput() {
+    return Dequantize<T>(ExtractVector<T>(output_), GetScale(output_),
+                         GetZeroPoint(output_));
+  }
+};
+
+class BoolConcatenationOpModel : public BaseConcatenationOpModel {
+ public:
+  using BaseConcatenationOpModel::BaseConcatenationOpModel;
+  void SetInput(int index, std::initializer_list<bool> data) {
+    PopulateTensor(index, data);
+  }
+  std::vector<bool> GetOutput() { return ExtractVector<bool>(output_); }
+};
+
+TEST(ConcatenationOpTest, ThreeDimensionalOneInputInt4) {
+  // INT4 values are packed 2 per byte.
+  // Shape {2, 1, 2} means 4 elements.
+  // Input: {1, 3, 4, 7}
+  // Packed:
+  // Byte 0: (1 & 0xF) | (3 << 4) = 0x31
+  // Byte 1: (4 & 0xF) | (7 << 4) = 0x74
+  ConcatenationOpModel<uint8_t> m0({TensorType_INT4, {2, 1, 2}}, /*axis=*/1,
+                                   /*num_inputs=*/1);
+  m0.SetInput(0, {0x31, 0x74});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({0x31, 0x74}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalOneInput) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 1, 2}}, /*axis=*/1,
+                                 /*num_inputs=*/1);
+  m0.SetInput(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 3, 4, 7}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalOneInputBFloat16) {
+  ConcatenationOpModel<Eigen::bfloat16> m({TensorType_BFLOAT16, {2, 1, 2}},
+                                          /*axis=*/1,
+                                          /*num_inputs=*/1);
+  m.SetInput(
+      0,
+      {static_cast<Eigen::bfloat16>(1.0f), static_cast<Eigen::bfloat16>(3.0f),
+       static_cast<Eigen::bfloat16>(4.0f), static_cast<Eigen::bfloat16>(7.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 4, 7}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalOneInputFloat16) {
+  ConcatenationOpModel<half> m({TensorType_FLOAT16, {2, 1, 2}},
+                               /*axis=*/1,
+                               /*num_inputs=*/1);
+  m.SetInput(0, {static_cast<half>(1.0f), static_cast<half>(3.0f),
+                 static_cast<half>(4.0f), static_cast<half>(7.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 4, 7}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalOneInputUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 1, 2}}, /*axis=*/1,
+                                    /*num_inputs=*/1);
+  m0.SetInput(0, {1, 3, 4, 7});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 3, 4, 7}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalOneInput) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 1, 2, 1, 3}},
+                                 /*axis=*/2,
+                                 /*num_inputs=*/1);
+  m0.SetInput(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f,
+                  11.0f, 12.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalOneInputUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 1, 2, 1, 3}},
+                                    /*axis=*/2,
+                                    /*num_inputs=*/1);
+  m0.SetInput(0, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInput) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 1, 2, 1, 3}},
+                                 /*axis=*/0,
+                                 /*num_inputs=*/2);
+  m0.SetInput(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f,
+                  11.0f, 12.0f});
+  m0.SetInput(1, {13.0f, 14.0f, 15.0f, 16.0f, 17.0f, 18.0f, 19.0f, 20.0f, 21.0f,
+                  22.0f, 23.0f, 24.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m0.GetOutput(),
+      ElementsAreArray({1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                        13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputBFloat16) {
+  ConcatenationOpModel<Eigen::bfloat16> m(
+      {TensorType_BFLOAT16, {2, 1, 2, 1, 3}},
+      /*axis=*/0,
+      /*num_inputs=*/2);
+  m.SetInput(
+      0,
+      {static_cast<Eigen::bfloat16>(1.0f), static_cast<Eigen::bfloat16>(2.0f),
+       static_cast<Eigen::bfloat16>(3.0f), static_cast<Eigen::bfloat16>(4.0f),
+       static_cast<Eigen::bfloat16>(5.0f), static_cast<Eigen::bfloat16>(6.0f),
+       static_cast<Eigen::bfloat16>(7.0f), static_cast<Eigen::bfloat16>(8.0f),
+       static_cast<Eigen::bfloat16>(9.0f), static_cast<Eigen::bfloat16>(10.0f),
+       static_cast<Eigen::bfloat16>(11.0f),
+       static_cast<Eigen::bfloat16>(12.0f)});
+  m.SetInput(
+      1,
+      {static_cast<Eigen::bfloat16>(13.0f), static_cast<Eigen::bfloat16>(14.0f),
+       static_cast<Eigen::bfloat16>(15.0f), Eigen::bfloat16{16.0f},
+       static_cast<Eigen::bfloat16>(17.0f), static_cast<Eigen::bfloat16>(18.0f),
+       static_cast<Eigen::bfloat16>(19.0f), static_cast<Eigen::bfloat16>(20.0f),
+       static_cast<Eigen::bfloat16>(21.0f), static_cast<Eigen::bfloat16>(22.0f),
+       static_cast<Eigen::bfloat16>(23.0f),
+       static_cast<Eigen::bfloat16>(24.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                        13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputFloat16) {
+  ConcatenationOpModel<half> m({TensorType_FLOAT16, {2, 1, 2, 1, 3}},
+                               /*axis=*/0,
+                               /*num_inputs=*/2);
+  m.SetInput(0, {static_cast<half>(1.0f), static_cast<half>(2.0f),
+                 static_cast<half>(3.0f), static_cast<half>(4.0f),
+                 static_cast<half>(5.0f), static_cast<half>(6.0f),
+                 static_cast<half>(7.0f), half{8.0f}, static_cast<half>(9.0f),
+                 static_cast<half>(10.0f), static_cast<half>(11.0f),
+                 static_cast<half>(12.0f)});
+  m.SetInput(1,
+             {static_cast<half>(13.0f), static_cast<half>(14.0f), half{15.0f},
+              static_cast<half>(16.0f), half{17.0f}, static_cast<half>(18.0f),
+              static_cast<half>(19.0f), static_cast<half>(20.0f),
+              static_cast<half>(21.0f), static_cast<half>(22.0f),
+              static_cast<half>(23.0f), static_cast<half>(24.0f)});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                        13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 1, 2, 1, 3}},
+                                    /*axis=*/0,
+                                    /*num_inputs=*/2);
+  m0.SetInput(0, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  m0.SetInput(1, {13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m0.GetOutput(),
+      ElementsAreArray({1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 12,
+                        13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputNegativeAxes) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 1, 2, 1, 3}},
+                                 /*axis=*/-2,
+                                 /*num_inputs=*/2);
+  m0.SetInput(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f,
+                  11.0f, 12.0f});
+  m0.SetInput(1, {13.0f, 14.0f, 15.0f, 16.0f, 17.0f, 18.0f, 19.0f, 20.0f, 21.0f,
+                  22.0f, 23.0f, 24.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 13, 14, 15, 4,  5,  6,  16, 17, 18,
+                                7, 8, 9, 19, 20, 21, 10, 11, 12, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputNegativeAxesUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 1, 2, 1, 3}},
+                                    /*axis=*/-2,
+                                    /*num_inputs=*/2);
+  m0.SetInput(0, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  m0.SetInput(1, {13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 13, 14, 15, 4,  5,  6,  16, 17, 18,
+                                7, 8, 9, 19, 20, 21, 10, 11, 12, 22, 23, 24}));
+}
+
+TEST(ConcatenationOpTest, FiveDimensionalTwoInputQuantizedUint8) {
+  QuantizedConcatenationOpModel m0(
+      {TensorType_UINT8, {2, 1, 2, 1, 3}, -12.7, 12.8},
+      /*axis=*/0,
+      /*num_inputs=*/2);
+
+  m0.SetInput<uint8_t>(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f,
+                           10.0f, 11.0f, 12.0f});
+  m0.SetInput<uint8_t>(1, {1.1f, 2.1f, 3.1f, 4.1f, 5.1f, 6.1f, 7.1f, 8.1f, 9.1f,
+                           10.1f, 11.1f, 12.1f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetDequantizedOutput<uint8_t>(),
+              ElementsAreArray(ArrayFloatNear({
+                  1.0f, 2.0f,  3.0f,  4.0f,  5.0f, 6.0f,  7.0f,  8.0f,
+                  9.0f, 10.0f, 11.0f, 12.0f, 1.1f, 2.1f,  3.1f,  4.1f,
+                  5.1f, 6.1f,  7.1f,  8.1f,  9.1f, 10.1f, 11.1f, 12.1f,
+              })));
+  EXPECT_THAT(
+      m0.GetOutput<uint8_t>(),
+      ElementsAreArray({
+          137, 147, 157, 167, 177, 187, 197, 207, 217, 227, 237, 247, 138,  //
+          148, 158, 168, 178, 188, 198, 208, 218, 228, 238, 248,
+      }));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalTwoInputsDifferentShapes) {
+  ConcatenationOpModel<float> m0(
+      {{TensorType_FLOAT32, {2, 1, 2}}, {TensorType_FLOAT32, {2, 3, 2}}},
+      /*axis=*/1, /*num_inputs=*/2, TensorType_FLOAT32);
+  m0.SetInput(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  m0.SetInput(1, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0, 7.0f, 8.0f, 9.0f, 10.0f,
+                  11.0f, 12.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 3, 1, 2, 3, 4, 5, 6, 4, 7, 7,
+                                                8, 9, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalTwoInputsDifferentShapesInt4) {
+  // Input 0: {2, 1, 2}, 4 elements -> {1, 3, 4, 7}
+  // Packed: 0x31, 0x74
+  // Input 1: {2, 3, 2}, 12 elements -> {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
+  // Packed: 0x21, 0x43, 0x65, 0x87, 0xA9, 0xCB
+  //
+  // Output: {2, 4, 2} (axis=1 concat), 16 elements
+  // Expected Output (logical):
+  // Row 0 (concat input 0 row 0 and input 1 row 0):
+  // {1, 3} (from in0) + {1, 2, 3, 4, 5, 6} (from in1)
+  // -> {1, 3, 1, 2, 3, 4, 5, 6}
+  // Packed: 0x31, 0x21, 0x43, 0x65
+  // Row 1 (concat input 0 row 1 and input 1 row 1):
+  // {4, 7} (from in0) + {7, 8, 9, 10, 11, 12} (from in1)
+  // -> {4, 7, 7, 8, 9, 10, 11, 12}
+  // Packed: 0x74, 0x87, 0xA9, 0xCB
+  //
+  // Total Packed Output: 0x31, 0x21, 0x43, 0x65, 0x74, 0x87, 0xA9, 0xCB
+
+  ConcatenationOpModel<uint8_t> m0(
+      {{TensorType_INT4, {2, 1, 2}}, {TensorType_INT4, {2, 3, 2}}},
+      /*axis=*/1, /*num_inputs=*/2, TensorType_INT4);
+  m0.SetInput(0, {0x31, 0x74});
+  m0.SetInput(1, {0x21, 0x43, 0x65, 0x87, 0xA9, 0xCB});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({0x31, 0x21, 0x43, 0x65, 0x74,
+                                                0x87, 0xA9, 0xCB}));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalTwoInputsDifferentShapesUInt32) {
+  ConcatenationOpModel<uint32_t> m0(
+      {{TensorType_UINT32, {2, 1, 2}}, {TensorType_UINT32, {2, 3, 2}}},
+      /*axis=*/1, /*num_inputs=*/2, TensorType_UINT32);
+  m0.SetInput(0, {1, 3, 4, 7});
+  m0.SetInput(1, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 3, 1, 2, 3, 4, 5, 6, 4, 7, 7,
+                                                8, 9, 10, 11, 12}));
+}
+
+#if GTEST_HAS_DEATH_TEST
+TEST(ConcatenationOpTest, ThreeDimensionalTwoInputsDifferentShapesWrongAxis) {
+  EXPECT_DEATH(
+      ConcatenationOpModel<float> m0(
+          {{TensorType_FLOAT32, {2, 1, 2}}, {TensorType_FLOAT32, {2, 3, 2}}},
+          /*axis=*/0, /*num_inputs=*/2, TensorType_FLOAT32),
+      "Cannot allocate tensors");
+}
+
+TEST(ConcatenationOpTest,
+     ThreeDimensionalTwoInputsDifferentShapesWrongAxisUInt32) {
+  EXPECT_DEATH(
+      ConcatenationOpModel<uint32_t> m0(
+          {{TensorType_UINT32, {2, 1, 2}}, {TensorType_UINT32, {2, 3, 2}}},
+          /*axis=*/0, /*num_inputs=*/2, TensorType_UINT32),
+      "Cannot allocate tensors");
+}
+#endif
+
+TEST(ConcatenationOpTest, OneTrivialInput) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {1}}, /*axis=*/0,
+                                 /*num_inputs=*/1);
+  m0.SetInput(0, {5.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ::testing::ElementsAre(5));
+}
+
+TEST(ConcatenationOpTest, OneTrivialInputUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {1}}, /*axis=*/0,
+                                    /*num_inputs=*/1);
+  m0.SetInput(0, {5});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ::testing::ElementsAre(5));
+}
+
+TEST(ConcatenationOpTest, TwoDimensionalOneInput) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 3}}, /*axis=*/0,
+                                 /*num_inputs=*/1);
+  m0.SetInput(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TEST(ConcatenationOpTest, TwoDimensionalOneInputUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 3}}, /*axis=*/0,
+                                    /*num_inputs=*/1);
+  m0.SetInput(0, {1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TEST(ConcatenationOpTest, TwoInputsTwoAxesNegativeAxes) {
+  // We will concatenate two tensors along different dimensions.
+  auto tensor0 = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+  auto tensor1 = {7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f};
+
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 3}}, /*axis=*/0,
+                                 /*num_inputs=*/2);
+  m0.SetInput(0, tensor0);
+  m0.SetInput(1, tensor1);
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  ConcatenationOpModel<float> m0_negative({TensorType_FLOAT32, {2, 3}},
+                                          /*axis=*/-2,
+                                          /*num_inputs=*/2);
+  m0_negative.SetInput(0, tensor0);
+  m0_negative.SetInput(1, tensor1);
+  ASSERT_EQ(m0_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0_negative.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  ConcatenationOpModel<float> m1({TensorType_FLOAT32, {2, 3}}, /*axis=*/1,
+                                 /*num_inputs=*/2);
+  m1.SetInput(0, tensor0);
+  m1.SetInput(1, tensor1);
+  ASSERT_EQ(m1.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1.GetOutput(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+
+  ConcatenationOpModel<float> m1_negative({TensorType_FLOAT32, {2, 3}},
+                                          /*axis=*/-1,
+                                          /*num_inputs=*/2);
+  m1_negative.SetInput(0, tensor0);
+  m1_negative.SetInput(1, tensor1);
+  ASSERT_EQ(m1_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1_negative.GetOutput(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, TwoInputsTwoAxesNegativeAxesUInt32) {
+  // We will concatenate two tensors along different dimensions.
+  std::initializer_list<uint32_t> tensor0 = {1, 2, 3, 4, 5, 6};
+  std::initializer_list<uint32_t> tensor1 = {7, 8, 9, 10, 11, 12};
+
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 3}}, /*axis=*/0,
+                                    /*num_inputs=*/2);
+  m0.SetInput(0, tensor0);
+  m0.SetInput(1, tensor1);
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  ConcatenationOpModel<uint32_t> m0_negative({TensorType_UINT32, {2, 3}},
+                                             /*axis=*/-2,
+                                             /*num_inputs=*/2);
+  m0_negative.SetInput(0, tensor0);
+  m0_negative.SetInput(1, tensor1);
+  ASSERT_EQ(m0_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0_negative.GetOutput(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  ConcatenationOpModel<uint32_t> m1({TensorType_UINT32, {2, 3}}, /*axis=*/1,
+                                    /*num_inputs=*/2);
+  m1.SetInput(0, tensor0);
+  m1.SetInput(1, tensor1);
+  ASSERT_EQ(m1.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1.GetOutput(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+
+  ConcatenationOpModel<uint32_t> m1_negative({TensorType_UINT32, {2, 3}},
+                                             /*axis=*/-1,
+                                             /*num_inputs=*/2);
+  m1_negative.SetInput(0, tensor0);
+  m1_negative.SetInput(1, tensor1);
+  ASSERT_EQ(m1_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1_negative.GetOutput(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, FourInputs) {
+  ConcatenationOpModel<float> m0({TensorType_FLOAT32, {2, 1, 2}}, /*axis=*/2,
+                                 /*num_inputs=*/4);
+  m0.SetInput(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  m0.SetInput(1, {1.1f, 3.1f, 4.1f, 7.1f});
+  m0.SetInput(2, {1.2f, 3.2f, 4.2f, 7.2f});
+  m0.SetInput(3, {1.3f, 3.3f, 4.3f, 7.3f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(),
+              Pointwise(FloatingPointEq(),
+                        {
+                            1.0f, 3.0f, 1.1f, 3.1f, 1.2f, 3.2f, 1.3f, 3.3f,  //
+                            4.0f, 7.0f, 4.1f, 7.1f, 4.2f, 7.2f, 4.3f, 7.3f,  //
+                        }));
+}
+
+TEST(ConcatenationOpTest, FourInputsUInt32) {
+  ConcatenationOpModel<uint32_t> m0({TensorType_UINT32, {2, 1, 2}}, /*axis=*/2,
+                                    /*num_inputs=*/4);
+  m0.SetInput(0, {1, 3, 4, 7});
+  m0.SetInput(1, {1, 3, 4, 7});
+  m0.SetInput(2, {1, 3, 4, 7});
+  m0.SetInput(3, {1, 3, 4, 7});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({
+                                  1, 3, 1, 3, 1, 3, 1, 3,  //
+                                  4, 7, 4, 7, 4, 7, 4, 7,  //
+                              }));
+}
+
+TEST(ConcatenationOpTest, FourInputsQuantizedUint8) {
+  QuantizedConcatenationOpModel m0({TensorType_UINT8, {2, 1, 2}, -12.7, 12.8},
+                                   /*axis=*/2,
+                                   /*num_inputs=*/4);
+
+  m0.SetInput<uint8_t>(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  m0.SetInput<uint8_t>(1, {1.1f, 3.1f, 4.1f, 7.1f});
+  m0.SetInput<uint8_t>(2, {1.2f, 3.2f, 4.2f, 7.2f});
+  m0.SetInput<uint8_t>(3, {1.3f, 3.3f, 4.3f, 7.3f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetDequantizedOutput<uint8_t>(),
+              ElementsAreArray(ArrayFloatNear({
+                  1.0f, 3.0f, 1.1f, 3.1f, 1.2f, 3.2f, 1.3f, 3.3f,  //
+                  4.0f, 7.0f, 4.1f, 7.1f, 4.2f, 7.2f, 4.3f, 7.3f,  //
+              })));
+  EXPECT_THAT(m0.GetOutput<uint8_t>(),
+              ElementsAreArray({
+                  137, 157, 138, 158, 139, 159, 140, 160,  //
+                  167, 197, 168, 198, 169, 199, 170, 200,  //
+              }));
+}
+
+template <typename Type>
+struct ConcatenationOpTestTyped : public testing::Test {
+  using TestType = Type;
+
+  enum TensorType tensor_type =
+      (std::is_same<Type, int16_t>::value ? TensorType_INT16 : TensorType_INT8);
+};
+
+using TestTypes = testing::Types<int8_t, int16_t>;
+TYPED_TEST_CASE(ConcatenationOpTestTyped, TestTypes);
+
+TYPED_TEST(ConcatenationOpTestTyped, FourInputsQuantizedInt8) {
+  using TestType = typename TestFixture::TestType;
+
+  const float kMin = -1;
+  const float kMax =
+      std::numeric_limits<TestType>::max() /
+      static_cast<float>(std::numeric_limits<TestType>::max() + 1);
+
+  QuantizedConcatenationOpModel m0(
+      {TestFixture::tensor_type, {2, 1, 2}, 12.8f * kMin, 12.8f * kMax},
+      /*axis=*/2,
+      /*num_inputs=*/4);
+
+  m0.SetInput<TestType>(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  m0.SetInput<TestType>(1, {1.1f, 3.1f, 4.1f, 7.1f});
+  m0.SetInput<TestType>(2, {1.2f, 3.2f, 4.2f, 7.2f});
+  m0.SetInput<TestType>(3, {1.3f, 3.3f, 4.3f, 7.3f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetDequantizedOutput<TestType>(),
+              ElementsAreArray(ArrayFloatNear({
+                  1, 3, 1.1, 3.1, 1.2, 3.2, 1.3, 3.3,  //
+                  4, 7, 4.1, 7.1, 4.2, 7.2, 4.3, 7.3   //
+              })));
+}
+
+TEST(ConcatenationOpTest, FourInputsQuantizedMixedRange) {
+  QuantizedConcatenationOpModel m0({{TensorType_UINT8, {2, 1, 2}, -10.7, 10.8},
+                                    {TensorType_UINT8, {2, 1, 2}, 0, 12.8},
+                                    {TensorType_UINT8, {2, 1, 2}, -11, 11.8},
+                                    {TensorType_UINT8, {2, 1, 2}, 0, 7.4}},
+                                   /*axis=*/2, /*num_inputs=*/4,
+                                   {TensorType_UINT8, {2, 1, 2}, -12.7, 12.8});
+
+  m0.SetInput<uint8_t>(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  m0.SetInput<uint8_t>(1, {1.1f, 3.1f, 4.1f, 7.1f});
+  m0.SetInput<uint8_t>(2, {1.2f, 3.2f, 4.2f, 7.2f});
+  m0.SetInput<uint8_t>(3, {1.3f, 3.3f, 4.3f, 7.3f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetDequantizedOutput<uint8_t>(),
+              ElementsAreArray(ArrayFloatNear({
+                  1.0f, 3.0f, 1.1f, 3.1f, 1.2f, 3.2f, 1.3f, 3.3f,  //
+                  4.0f, 7.0f, 4.1f, 7.1f, 4.2f, 7.2f, 4.3f, 7.3f,  //
+              })));
+  EXPECT_THAT(m0.GetOutput<uint8_t>(),
+              ElementsAreArray({
+                  137, 157, 138, 158, 139, 159, 140, 160,  //
+                  167, 197, 168, 198, 169, 199, 170, 200,  //
+              }));
+}
+
+TEST(ConcatenationOpTest, FourInputsQuantizedMixedRangeClampingLogic) {
+  QuantizedConcatenationOpModel m0({{TensorType_UINT8, {2, 1, 2}, -10.7, 10.8},
+                                    {TensorType_UINT8, {2, 1, 2}, 0, 12.8},
+                                    {TensorType_UINT8, {2, 1, 2}, -11, 11.8},
+                                    {TensorType_UINT8, {2, 1, 2}, 0, 7.4}},
+                                   /*axis=*/2, /*num_inputs=*/4,
+                                   {TensorType_UINT8, {2, 1, 2}, -1., 1.});
+
+  m0.SetInput<uint8_t>(0, {1.0f, -3.0f, -4.0f, -7.0f});
+  m0.SetInput<uint8_t>(1, {1.1f, 3.1f, 4.1f, 7.1f});
+  m0.SetInput<uint8_t>(2, {1.2f, -3.2f, -4.2f, 7.2f});
+  m0.SetInput<uint8_t>(3, {1.3f, 3.3f, 4.3f, 7.3f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetDequantizedOutput<uint8_t>(),
+              ElementsAreArray(ArrayFloatNear(
+                  {
+                      1.0f, -1.0f, 1.0f, 1.0f, 1.0f, -1.0f, 1.0f, 1.0f,   //
+                      -1.0f, -1.0f, 1.0f, 1.0f, -1.0f, 1.0f, 1.0f, 1.0f,  //
+                  },
+                  4e-3)));
+  EXPECT_THAT(m0.GetOutput<uint8_t>(),
+              ElementsAreArray({
+                  255, 0, 255, 255, 255, 0, 255, 255,  //
+                  0, 0, 255, 255, 0, 255, 255, 255,    //
+              }));
+}
+
+TEST(ConcatenationOpTest, ThreeDimensionalNonQuantizedOneInput) {
+  QuantizedConcatenationOpModel m0(
+      {TensorType_UINT8, {2, 1, 2}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/1,
+      /*num_inputs=*/1);
+  m0.SetInput<uint8_t>(0, {1.0f, 3.0f, 4.0f, 7.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput<uint8_t>(),
+              ElementsAreArray(ArrayFloatNear({1.0f, 3.0f, 4.0f, 7.0f})));
+}
+
+TEST(ConcatenationOpTest, OneTrivialNonQuantizedInput) {
+  QuantizedConcatenationOpModel m0(
+      {TensorType_UINT8, {1}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/0,
+      /*num_inputs=*/1);
+  m0.SetInput<uint8_t>(0, {5.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput<uint8_t>(), ::testing::ElementsAre(5));
+}
+
+TEST(ConcatenationOpTest, TwoDimensionalNonQuantizedOneInput) {
+  QuantizedConcatenationOpModel m0(
+      {TensorType_UINT8, {2, 3}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/0,
+      /*num_inputs=*/1);
+  m0.SetInput<uint8_t>(0, {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput<uint8_t>(), ElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TEST(ConcatenationOpTest, TwoInputsTwoAxesNegativeAxesNonQuantized) {
+  // We will concatenate two tensors along different dimensions.
+  auto tensor0 = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+  auto tensor1 = {7.0f, 8.0f, 9.0f, 10.0f, 11.0f, 12.0f};
+
+  QuantizedConcatenationOpModel m0(
+      {TensorType_UINT8, {2, 3}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/0,
+      /*num_inputs=*/2);
+  m0.SetInput<uint8_t>(0, tensor0);
+  m0.SetInput<uint8_t>(1, tensor1);
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput<uint8_t>(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  QuantizedConcatenationOpModel m0_negative(
+      {TensorType_UINT8, {2, 3}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/-2,
+      /*num_inputs=*/2);
+  m0_negative.SetInput<uint8_t>(0, tensor0);
+  m0_negative.SetInput<uint8_t>(1, tensor1);
+  ASSERT_EQ(m0_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0_negative.GetOutput<uint8_t>(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+
+  QuantizedConcatenationOpModel m1(
+      {TensorType_UINT8, {2, 3}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/1,
+      /*num_inputs=*/2);
+  m1.SetInput<uint8_t>(0, tensor0);
+  m1.SetInput<uint8_t>(1, tensor1);
+  ASSERT_EQ(m1.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1.GetOutput<uint8_t>(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+
+  QuantizedConcatenationOpModel m1_negative(
+      {TensorType_UINT8, {2, 3}, 0, std::numeric_limits<uint8_t>::max()},
+      /*axis=*/-1,
+      /*num_inputs=*/2);
+  m1_negative.SetInput<uint8_t>(0, tensor0);
+  m1_negative.SetInput<uint8_t>(1, tensor1);
+  ASSERT_EQ(m1_negative.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m1_negative.GetOutput<uint8_t>(),
+              ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12}));
+}
+
+TEST(ConcatenationOpTest, BoolTypeOneInput) {
+  BoolConcatenationOpModel m0({TensorType_BOOL, {2, 1, 2}}, /*axis=*/1,
+                              /*num_inputs=*/1);
+  m0.SetInput(0, {true, false, false, true});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray({true, false, false, true}));
+}
+
+TEST(ConcatenationOpTest, BoolTypeTwoInputs) {
+  BoolConcatenationOpModel m0(
+      {{TensorType_BOOL, {2, 1, 2}}, {TensorType_BOOL, {2, 3, 2}}},
+      /*axis=*/1, /*num_inputs=*/2, TensorType_BOOL);
+  m0.SetInput(0, {false, false, false, false});
+  m0.SetInput(1, {true, true, true, true, true, true, true, true, true, true,
+                  true, true});
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m0.GetOutput(),
+      ElementsAreArray({false, false, true, true, true, true, true, true, false,
+                        false, true, true, true, true, true, true}));
+}
+
+enum class TestInputType {
+  kPersistentRo = 0,
+  kOnePersistentRo = 1,
+  kDefault = 2,
+};
+
+struct PersistentTestCase {
+  TestInputType test_type;
+  TensorType tensor_type;
+  bool is_quantized = false;
+};
+
+template <typename T>
+class PersistentConcatenationOpModel : public SingleOpModel {
+ public:
+  PersistentConcatenationOpModel(const std::vector<TensorData>& input_template,
+                                 int axis, const TensorData& output_template,
+                                 PersistentTestCase test_case,
+                                 std::vector<std::vector<T>> input_data_list)
+      : input_data_list_(input_data_list), test_case_(test_case) {
+    const int num_inputs = input_data_list.size();
+    std::vector<std::vector<int>> all_input_shapes;
+    CHECK_EQ(input_template.size(), num_inputs);
+    for (int i = 0; i < num_inputs; ++i) {
+      int id;
+      all_input_shapes.push_back(input_template[i].shape);
+      id = AddInput(input_template[i]);
+      concat_inputs_.push_back(id);
+    }
+    output_ = AddOutput(output_template);
+    SetBuiltinOp(
+        BuiltinOperator_CONCATENATION, BuiltinOptions_ConcatenationOptions,
+        CreateConcatenationOptions(builder_, axis, ActivationFunctionType_NONE)
+            .Union());
+    BuildInterpreter(all_input_shapes, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true,
+                     /*allocate_and_delegate=*/false);
+
+    int num_persistent_inputs = 0;
+    if (test_case_.test_type == TestInputType::kPersistentRo) {
+      num_persistent_inputs = num_inputs;
+    } else if (test_case_.test_type == TestInputType::kOnePersistentRo) {
+      num_persistent_inputs = 1;
+    }
+
+    for (int i = 0; i < num_persistent_inputs; ++i) {
+      interpreter_->tensor(concat_inputs_[i])->allocation_type =
+          kTfLitePersistentRo;
+      std::vector<T>& input_data = input_data_list[i];
+      interpreter_->ResizeInputTensorStrict(concat_inputs_[i],
+                                            input_template[i].shape);
+      if (test_case.is_quantized) {
+        QuantizeAndPopulate<int8_t>(concat_inputs_[i], FloatVector(input_data));
+      } else {
+        PopulateTensor(concat_inputs_[i], input_data);
+      }
+    }
+    AllocateAndDelegate(true);
+  }
+
+  std::vector<float> FloatVector(std::vector<T> data) {
+    std::vector<float> ret;
+    for (T t : data) {
+      ret.push_back(static_cast<float>(t));
+    }
+    return ret;
+  }
+
+  void PopulateInputTensors() {
+    int start = -1;
+    if (test_case_.test_type == TestInputType::kDefault) {
+      start = 0;
+    } else if (test_case_.test_type == TestInputType::kOnePersistentRo) {
+      start = 1;
+    }
+    if (start < 0) {
+      return;
+    }
+    for (int i = start; i < input_data_list_.size(); ++i) {
+      if (test_case_.is_quantized) {
+        QuantizeAndPopulate<int8_t>(concat_inputs_[i],
+                                    FloatVector(input_data_list_[i]));
+      } else {
+        std::vector<T> v(input_data_list_[i]);
+        PopulateTensor(concat_inputs_[i], v);
+      }
+    }
+  }
+
+  bool IsPersistentOutput() {
+    const TfLiteTensor* tensor = interpreter_->tensor(output_);
+    return tensor->allocation_type == kTfLitePersistentRo;
+  }
+
+  std::vector<float> GetOutput() {
+    if (test_case_.is_quantized) {
+      return Dequantize<int8_t>(ExtractVector<int8_t>(output_),
+                                GetScale(output_), GetZeroPoint(output_));
+    }
+    return FloatVector(ExtractVector<T>(output_));
+  }
+
+ protected:
+  int output_;
+  std::vector<std::vector<T>> input_data_list_;
+  PersistentTestCase test_case_;
+  std::vector<int> concat_inputs_;
+};
+
+template <typename T>
+class ConcatenationOpPersistentModelTest : public ::testing::Test {
+ public:
+  static std::vector<PersistentTestCase> Range(bool is_quantized = false) {
+    TensorType tensor_type = TensorType_FLOAT32;
+    if (std::is_same<T, int32_t>::value) {
+      tensor_type = TensorType_INT32;
+    }
+    if (std::is_same<T, uint32_t>::value) {
+      tensor_type = TensorType_UINT32;
+    }
+    if (is_quantized) {
+      tensor_type = TensorType_INT8;
+    }
+    return {{TestInputType::kDefault, tensor_type, is_quantized},
+            {TestInputType::kPersistentRo, tensor_type, is_quantized}};
+  }
+};
+
+using DataTypes = ::testing::Types<float, int32_t, uint32_t>;
+TYPED_TEST_SUITE(ConcatenationOpPersistentModelTest, DataTypes);
+
+TYPED_TEST(ConcatenationOpPersistentModelTest, PersistentTest) {
+  for (PersistentTestCase test_case :
+       ConcatenationOpPersistentModelTest<TypeParam>::Range()) {
+    std::vector<std::vector<TypeParam>> input_data_lists = {
+        {1, 2, 3, 4, 5, 6}, {7, 8, 9, 10, 11, 12}};
+    std::vector<TensorData> input_template = {{test_case.tensor_type, {2, 3}},
+                                              {test_case.tensor_type, {2, 3}}};
+    TensorData output_template = {test_case.tensor_type, {4, 3}};
+    PersistentConcatenationOpModel<TypeParam> m0(input_template, /*axis=*/0,
+                                                 output_template, test_case,
+                                                 input_data_lists);
+    m0.PopulateInputTensors();
+    ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+    if (m0.GetNumberOfAppliedDelegates() == 0) {
+      ASSERT_EQ(m0.IsPersistentOutput(),
+                test_case.test_type == TestInputType::kPersistentRo);
+    }
+    EXPECT_THAT(
+        m0.GetOutput(),
+        ElementsAreArray(ArrayFloatNear(
+            {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0})));
+  }
+}
+
+TYPED_TEST(ConcatenationOpPersistentModelTest, PersistentScalarTest) {
+  PersistentTestCase test_case{TestInputType::kPersistentRo,
+                               GetTensorType<TypeParam>(), false};
+  std::vector<std::vector<TypeParam>> input_data_lists = {{1}, {7}};
+  std::vector<TensorData> input_template = {{GetTensorType<TypeParam>(), {}},
+                                            {GetTensorType<TypeParam>(), {}}};
+  TensorData output_template = {GetTensorType<TypeParam>(), {2}};
+  PersistentConcatenationOpModel<TypeParam> m0(
+      input_template, /*axis=*/0, output_template, test_case, input_data_lists);
+  m0.PopulateInputTensors();
+  ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+  ASSERT_EQ(m0.IsPersistentOutput(),
+            test_case.test_type == TestInputType::kPersistentRo);
+  EXPECT_THAT(m0.GetOutput(), ElementsAreArray(ArrayFloatNear({1.0, 7.0})));
+}
+
+TYPED_TEST(ConcatenationOpPersistentModelTest, QuantizedPersistentTest) {
+  const bool is_quantized = true;
+  for (PersistentTestCase test_case :
+       ConcatenationOpPersistentModelTest<TypeParam>::Range(is_quantized)) {
+    std::vector<std::vector<TypeParam>> input_data_lists = {
+        {1, 2, 3, 4, 5, 6}, {7, 8, 9, 10, 11, 12}};
+    float scale = 12.0 / 255.0;
+    int zero_point = -128;
+    std::vector<TensorData> input_template = {
+        {test_case.tensor_type, {2, 3}, 0.0, 12.0, scale, zero_point},
+        {test_case.tensor_type, {2, 3}, 0.0, 12.0, scale, zero_point},
+    };
+    TensorData output_template = {
+        test_case.tensor_type, {4, 3}, 0.0, 12.0, scale, zero_point};
+    PersistentConcatenationOpModel<TypeParam> m0(input_template, /*axis=*/0,
+                                                 output_template, test_case,
+                                                 input_data_lists);
+    m0.PopulateInputTensors();
+    ASSERT_EQ(m0.Invoke(), kTfLiteOk);
+    if (m0.GetNumberOfAppliedDelegates() == 0) {
+      ASSERT_EQ(m0.IsPersistentOutput(),
+                test_case.test_type == TestInputType::kPersistentRo);
+    }
+    EXPECT_THAT(
+        m0.GetOutput(),
+        ElementsAreArray(ArrayFloatNear(
+            {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0},
+            1e-1)));
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/control_flow_common.cc

@@ -0,0 +1,33 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/control_flow_common.h"
+
+#include <algorithm>
+#include <vector>
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+
+int OutputIsInput(int output_idx, const std::vector<int>& subgraph_inputs) {
+  auto e =
+      std::find(subgraph_inputs.begin(), subgraph_inputs.end(), output_idx);
+  return (e != subgraph_inputs.end()) ? (e - subgraph_inputs.begin()) : -1;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/control_flow_common.h

@@ -0,0 +1,172 @@
+/* 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_CONTROL_FLOW_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_CONTROL_FLOW_COMMON_H_
+
+#include <vector>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/subgraph.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+// Propagate tensor shapes and types from `src_tensor_indices` in `src_subgraph`
+// to `dst_tensor_indices` in `dst_subgraph`.
+//
+// When `resize_subgraph_inputs` is true, the function calls subgraphs's
+// `ResizeInputTensor` function, and it may trigger the memory planner to
+// reallocate memory.
+// When `resize_subgraph_inputs` is false, it implies `context` belongs to
+// `dst_subgraph`. The function calls `context->ResizeTensor`. This happens
+// when resizing `While` op's outputs.
+template <typename SrcVector, typename DstVector>
+TfLiteStatus CopyTensorsShapeAndType(TfLiteContext* context,
+                                     Subgraph* src_subgraph,
+                                     const SrcVector& src_tensor_indices,
+                                     Subgraph* dst_subgraph,
+                                     const DstVector& dst_tensor_indices,
+                                     bool resize_subgraph_inputs) {
+  TF_LITE_ENSURE_EQ(context, src_tensor_indices.size(),
+                    dst_tensor_indices.size());
+  for (int i = 0; i < src_tensor_indices.size(); ++i) {
+    // Skip copying unused destination tensors.
+    if (dst_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+
+    const TfLiteTensor* src_tensor =
+        src_subgraph->tensor(src_tensor_indices[i]);
+
+    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
+    if (resize_subgraph_inputs) {
+      dst_subgraph->ResizeInputTensor(dst_tensor_indices[i],
+                                      src_tensor->dims->data,
+                                      src_tensor->dims->size);
+    } else {
+      TF_LITE_ENSURE_OK(
+          context, context->ResizeTensor(context, dst_tensor,
+                                         TfLiteIntArrayCopy(src_tensor->dims)));
+    }
+    dst_tensor->type = src_tensor->type;
+  }
+  return kTfLiteOk;
+}
+
+// Copy the tensors data from tensors `src_tensor_indices` in `src_subgraph`
+// to `dst_tensor_indices` in `dst_subgraph`.
+template <typename SrcVector, typename DstVector>
+TfLiteStatus CopyTensorsData(TfLiteContext* context, Subgraph* src_subgraph,
+                             const SrcVector& src_tensor_indices,
+                             Subgraph* dst_subgraph,
+                             const DstVector& dst_tensor_indices) {
+  TF_LITE_ENSURE_EQ(context, src_tensor_indices.size(),
+                    dst_tensor_indices.size());
+  for (int i = 0; i < src_tensor_indices.size(); ++i) {
+    // Skip copying unused destination tensors.
+    if (dst_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+
+    const TfLiteTensor* src_tensor =
+        src_subgraph->tensor(src_tensor_indices[i]);
+    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
+    if (IsDynamicTensor(dst_tensor)) {
+      TfLiteTensorRealloc(src_tensor->bytes, dst_tensor);
+    }
+    TF_LITE_ENSURE_OK(context, TfLiteTensorCopy(src_tensor, dst_tensor));
+  }
+  return kTfLiteOk;
+}
+
+// Propagate tensor shapes and types from `src_tensor_indices` in `src_subgraph`
+// to `dst_tensor_indices` in `dst_subgraph` and copy data deeply.
+template <typename SrcVector, typename DstVector>
+TfLiteStatus DeepCopyTensorsShapeTypeData(
+    TfLiteContext* context, TfLiteNode* node, Subgraph* src_subgraph,
+    const SrcVector& src_tensor_indices, Subgraph* dst_subgraph,
+    const DstVector& dst_tensor_indices, bool body_has_dynamic_output_tensors) {
+  if (body_has_dynamic_output_tensors) {
+    Subgraph* this_subgraph = reinterpret_cast<Subgraph*>(context->impl_);
+    bool resize_subgraph_inputs = (dst_subgraph != this_subgraph);
+    TF_LITE_ENSURE_OK(
+        context, CopyTensorsShapeAndType(
+                     context, src_subgraph, src_tensor_indices, dst_subgraph,
+                     dst_tensor_indices, resize_subgraph_inputs));
+    if (resize_subgraph_inputs) {
+      TF_LITE_ENSURE_OK(context, dst_subgraph->AllocateTensors());
+    }
+  }
+  TF_LITE_ENSURE_OK(context,
+                    CopyTensorsData(context, src_subgraph, src_tensor_indices,
+                                    dst_subgraph, dst_tensor_indices));
+  return kTfLiteOk;
+}
+
+template <typename SrcVector, typename DstVector>
+TfLiteStatus DeepOrShallowCopyTensorsShapeTypeData(
+    TfLiteContext* context, TfLiteNode* node, Subgraph* src_subgraph,
+    const SrcVector& src_tensor_indices, Subgraph* dst_subgraph,
+    const DstVector& dst_tensor_indices) {
+  // Resize the destination subgraph inputs.
+  for (int i = 0; i < src_tensor_indices.size(); ++i) {
+    // Skip copying unused destination tensors.
+    if (dst_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+    if (src_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+
+    const TfLiteTensor* src_tensor =
+        src_subgraph->tensor(src_tensor_indices[i]);
+    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
+    std::vector<int> dims(src_tensor->dims->data,
+                          src_tensor->dims->data + src_tensor->dims->size);
+    dst_subgraph->ResizeInputTensor(dst_tensor_indices[i], dims);
+    dst_tensor->type = src_tensor->type;
+    if (!IsResourceOrVariant(src_tensor)) {
+      dst_tensor->bytes = 0;  // Don't allocate memory with AllocateTensors().
+      dst_tensor->data.raw = nullptr;
+    }
+  }
+  TF_LITE_ENSURE_OK(context, dst_subgraph->AllocateTensors());
+  // Deep or shallow copy the data from src subgraph to dst.
+  for (int i = 0; i < src_tensor_indices.size(); ++i) {
+    // Skip copying unused destination tensors.
+    if (dst_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+    if (src_tensor_indices[i] == kTfLiteOptionalTensor) continue;
+
+    const TfLiteTensor* src_tensor =
+        src_subgraph->tensor(src_tensor_indices[i]);
+    TfLiteTensor* dst_tensor = dst_subgraph->tensor(dst_tensor_indices[i]);
+    if (IsResourceOrVariant(src_tensor)) {
+      TfLiteTensorRealloc(src_tensor->bytes, dst_tensor);
+      TF_LITE_ENSURE_OK(context, TfLiteTensorCopy(src_tensor, dst_tensor));
+    } else {
+      // Make a shallow copy of the data. This is only safe because the caller
+      // is expected to have previously set dst_tensor->allocation_type to
+      // kTfLiteCustom, to ensure the buffer is never double-freed later on.
+      TF_LITE_ENSURE_EQ(context, dst_tensor->allocation_type, kTfLiteCustom);
+      dst_tensor->bytes = src_tensor->bytes;
+      dst_tensor->data.raw = src_tensor->data.raw;
+    }
+  }
+  return kTfLiteOk;
+}
+
+// Returns the subgraph input tensor index if the given output is also an input.
+// Otherwise returns -1.
+int OutputIsInput(int output_idx, const std::vector<int>& subgraph_inputs);
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CONTROL_FLOW_COMMON_H_

tensorflow/lite/kernels/conv3d.cc

@@ -0,0 +1,337 @@
+/* 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/internal/reference/conv3d.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/padding.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace conv3d {
+
+enum KernelType {
+  kReference,
+  kGenericOptimized,
+};
+
+// Struct to carry data from Prepare to Eval.
+const int kTensorNotAllocated = -1;
+static constexpr size_t kMaxIm2colBufferSizeMobile = 1024 * 1024 * 1024;  // 1GB
+
+struct OpData {
+  Padding3DValues padding;
+  int im2col_tensor_id = kTensorNotAllocated;
+
+  bool need_im2col = false;
+
+  // Disable im2col if the temporary im2col tensor requires too much memory
+  // (i.e. >= kMaxIm2colBufferSizeMobile).
+  bool im2col_oversized = false;
+
+  int32_t im2col_index;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* opdata = new OpData;
+  return opdata;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete static_cast<OpData*>(buffer);
+}
+
+TfLiteStatus AllocateTemporaryTensorsIfRequired(
+    KernelType kernel_type, TfLiteContext* context, TfLiteNode* node,
+    OpData* opdata, TfLiteConv3DParams* params, const TfLiteTensor* filter,
+    size_t im2col_bytes) {
+  int temporaries_count = 0;
+  const bool need_dilated_im2col = params->dilation_width_factor != 1 ||
+                                   params->dilation_height_factor != 1 ||
+                                   params->dilation_depth_factor != 1;
+  const bool need_non_dilated_im2col =
+      params->stride_depth != 1 || params->stride_width != 1 ||
+      params->stride_height != 1 || filter->dims->data[2] != 1 ||
+      filter->dims->data[1] != 1 || filter->dims->data[0] != 1;
+
+  opdata->need_im2col = (kernel_type == kGenericOptimized) &&
+                        (need_dilated_im2col || need_non_dilated_im2col);
+
+  // On mobile platforms, the generic optimized kernel will not be used if the
+  // temporary im2col tensor requires too much memory.
+  if (IsMobilePlatform() && opdata->need_im2col &&
+      im2col_bytes >= kMaxIm2colBufferSizeMobile) {
+    opdata->need_im2col = false;
+    opdata->im2col_oversized = true;
+  }
+
+  if (opdata->need_im2col) {
+    if (opdata->im2col_tensor_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &opdata->im2col_tensor_id));
+    }
+    opdata->im2col_index = temporaries_count++;
+  }
+
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(temporaries_count);
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
+                     TfLiteNode* node) {
+  auto* params = static_cast<TfLiteConv3DParams*>(node->builtin_data);
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+
+  // Check number of inputs/outputs.
+  TF_LITE_ENSURE(context, node->inputs->size == 2 || node->inputs->size == 3);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));
+
+  // Check dimensionality of input, filter.
+  TF_LITE_ENSURE_EQ(context, input->dims->size, 5);
+  TF_LITE_ENSURE_EQ(context, filter->dims->size, 5);
+
+  // Check input channels matching filter.
+  TF_LITE_ENSURE_EQ(context, input->dims->data[4], filter->dims->data[3]);
+
+  // Check types.
+  TfLiteType input_type = input->type;
+  TF_LITE_ENSURE_TYPES_EQ(context, input_type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, filter->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input_type);
+
+  // Check bias.
+  const TfLiteTensor* bias = GetInput(context, node, 2);
+  if (bias) {
+    TF_LITE_ENSURE_TYPES_EQ(context, bias->type, input_type);
+    TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 4));
+  }
+
+  // Filter has shape of [filter_depth, filter_height, filter_width,
+  // in_channels, out_channels].
+  int batches = input->dims->data[0];
+  int channels_out = filter->dims->data[4];
+  int depth = input->dims->data[1];
+  int height = input->dims->data[2];
+  int width = input->dims->data[3];
+  int filter_depth = filter->dims->data[0];
+  int filter_height = filter->dims->data[1];
+  int filter_width = filter->dims->data[2];
+  int input_channel = filter->dims->data[3];
+
+  // Matching GetWindowedOutputSize in TensorFlow.
+  int out_width, out_height, out_depth;
+  opdata->padding = ComputePadding3DValues(
+      params->stride_height, params->stride_width, params->stride_depth,
+      params->dilation_height_factor, params->dilation_width_factor,
+      params->dilation_depth_factor, height, width, depth, filter_height,
+      filter_width, filter_depth, params->padding, &out_height, &out_width,
+      &out_depth);
+
+  std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> output_size(
+      TfLiteIntArrayCreate(5), TfLiteIntArrayFree);
+  output_size->data[0] = batches;
+  output_size->data[1] = out_depth;
+  output_size->data[2] = out_height;
+  output_size->data[3] = out_width;
+  output_size->data[4] = channels_out;
+  TF_LITE_ENSURE_OK(
+      context, context->ResizeTensor(context, output, output_size.release()));
+
+  // Allocate temporary tensors.
+  size_t input_type_size;
+  TF_LITE_ENSURE_STATUS(GetSizeOfType(context, input->type, &input_type_size));
+  size_t im2col_elements = 0;
+  TF_LITE_ENSURE_OK(
+      context,
+      CheckedShapeProduct(
+          context,
+          {batches, out_depth, out_height, out_width, input_channel,
+           filter_depth, filter_height, filter_width},
+          "Conv3D im2col tensor has too many elements.", im2col_elements));
+  size_t im2col_bytes = 0;
+  TF_LITE_ENSURE_MSG(context,
+                     MultiplyAndCheckOverflow(im2col_elements, input_type_size,
+                                              &im2col_bytes) == kTfLiteOk,
+                     "%s", "Conv3D im2col tensor is too large.");
+  TF_LITE_ENSURE_OK(context, AllocateTemporaryTensorsIfRequired(
+                                 kernel_type, context, node, opdata, params,
+                                 filter, im2col_bytes));
+
+  if (opdata->need_im2col) {
+    if (im2col_elements > std::numeric_limits<int32_t>::max()) {
+      TF_LITE_KERNEL_LOG(
+          context,
+          "Conv3D im2col elements (%zu) exceed the 32-bit integer limit.",
+          im2col_elements);
+      return kTfLiteError;
+    }
+    std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> im2col_size(
+        TfLiteIntArrayCreate(5), TfLiteIntArrayFree);
+    im2col_size->data[0] = batches;
+    im2col_size->data[1] = out_depth;
+    im2col_size->data[2] = out_height;
+    im2col_size->data[3] = out_width;
+    const RuntimeShape filter_shape = GetTensorShape(filter);
+    int im2col_depth = 0;
+    TF_LITE_ENSURE_MSG(
+        context, filter_shape.CheckedSizeToDimension(/*end=*/4, im2col_depth),
+        "%s", "Conv3D im2col tensor has too many channels.");
+    im2col_size->data[4] = im2col_depth;
+
+    TfLiteTensor* im2col;
+    node->temporaries->data[opdata->im2col_index] = opdata->im2col_tensor_id;
+    TfLiteStatus status =
+        GetTemporarySafe(context, node, opdata->im2col_index, &im2col);
+    if (status != kTfLiteOk) {
+      return status;
+    }
+    im2col->type = input->type;
+    im2col->allocation_type = kTfLiteArenaRw;
+    status = context->ResizeTensor(context, im2col, im2col_size.release());
+    TF_LITE_ENSURE_OK(context, status);
+  }
+
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  return Prepare(kernel_type, context, node);
+}
+
+TfLiteStatus EvalFloat(KernelType kernel_type, TfLiteContext* context,
+                       TfLiteNode* node, TfLiteConv3DParams* params,
+                       OpData* opdata, const TfLiteTensor* input,
+                       const TfLiteTensor* filter, const TfLiteTensor* bias,
+                       TfLiteTensor* im2col, TfLiteTensor* output) {
+  float output_activation_min, output_activation_max;
+  CalculateActivationRange(params->activation, &output_activation_min,
+                           &output_activation_max);
+
+  Conv3DParams runtime_params;
+  runtime_params.padding_values = opdata->padding;
+  runtime_params.stride_depth = params->stride_depth;
+  runtime_params.stride_height = params->stride_height;
+  runtime_params.stride_width = params->stride_width;
+  runtime_params.dilation_depth = params->dilation_depth_factor;
+  runtime_params.dilation_height = params->dilation_height_factor;
+  runtime_params.dilation_width = params->dilation_width_factor;
+  runtime_params.float_activation_min = output_activation_min;
+  runtime_params.float_activation_max = output_activation_max;
+  switch (kernel_type) {
+    case kReference: {
+      reference_ops::Conv3D(runtime_params, GetTensorShape(input),
+                            GetTensorData<float>(input), GetTensorShape(filter),
+                            GetTensorData<float>(filter), GetTensorShape(bias),
+                            GetTensorData<float>(bias), GetTensorShape(output),
+                            GetTensorData<float>(output));
+      return kTfLiteOk;
+    }
+    case kGenericOptimized: {
+      return optimized_ops::Conv3D(
+          runtime_params, GetTensorShape(input), GetTensorData<float>(input),
+          GetTensorShape(filter), GetTensorData<float>(filter),
+          GetTensorShape(bias), GetTensorData<float>(bias),
+          GetTensorShape(output), GetTensorData<float>(output),
+          GetTensorShape(im2col), GetTensorData<float>(im2col),
+          CpuBackendContext::GetFromContext(context));
+    }
+  }
+}
+
+TfLiteStatus Eval(KernelType kernel_type, TfLiteContext* context,
+                  TfLiteNode* node) {
+  auto* params = reinterpret_cast<TfLiteConv3DParams*>(node->builtin_data);
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));
+  const TfLiteTensor* bias = GetInput(context, node, 2);
+
+  TfLiteTensor* im2col = opdata->need_im2col
+                             ? &context->tensors[opdata->im2col_tensor_id]
+                             : nullptr;
+  // Fallback to reference execution path when im2col is needed but disabled.
+  if (opdata->im2col_oversized) {
+    kernel_type = kReference;
+  }
+
+  switch (input->type) {
+    case kTfLiteFloat32:
+      return EvalFloat(kernel_type, context, node, params, opdata, input,
+                       filter, bias, im2col, output);
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %s currently not supported.",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  return Eval(kernel_type, context, node);
+}
+
+}  // namespace conv3d
+
+TfLiteRegistration* Register_CONV_3D_REF() {
+  static TfLiteRegistration r = {conv3d::Init, conv3d::Free,
+                                 conv3d::Prepare<conv3d::kReference>,
+                                 conv3d::Eval<conv3d::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONV_3D_GENERIC_OPT() {
+  static TfLiteRegistration r = {conv3d::Init, conv3d::Free,
+                                 conv3d::Prepare<conv3d::kGenericOptimized>,
+                                 conv3d::Eval<conv3d::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONV_3D() {
+  return Register_CONV_3D_GENERIC_OPT();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/conv3d_test.cc

@@ -0,0 +1,340 @@
+/* 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 <initializer_list>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+namespace ops {
+namespace builtin {
+
+TfLiteRegistration* Register_CONV_3D_GENERIC_OPT();
+
+}  // namespace builtin
+}  // namespace ops
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+class Conv3dOpModel : public SingleOpModel {
+ public:
+  Conv3dOpModel(const TensorData& input, const TensorData& filter,
+                const TensorData& bias, const TensorData& output,
+                Padding padding = Padding_VALID, int32_t stride_depth = 1,
+                int32_t stride_width = 1, int32_t stride_height = 1,
+                ActivationFunctionType activation = ActivationFunctionType_NONE,
+                int32_t dilation_depth = 1, int32_t dilation_width = 1,
+                int32_t dilation_height = 1) {
+    input_ = AddInput(input);
+    filter_ = AddInput(filter);
+    bias_ = AddInput(bias);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    BuildInterpreter({GetShape(input_), GetShape(filter_), GetShape(bias_)});
+  }
+
+  Conv3dOpModel(const TensorData& input, const TensorData& filter,
+                const TensorData& output, Padding padding = Padding_VALID,
+                int32_t stride_depth = 1, int32_t stride_width = 1,
+                int32_t stride_height = 1,
+                ActivationFunctionType activation = ActivationFunctionType_NONE,
+                int32_t dilation_depth = 1, int32_t dilation_width = 1,
+                int32_t dilation_height = 1) {
+    input_ = AddInput(input);
+    filter_ = AddInput(filter);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    BuildInterpreter({GetShape(input_), GetShape(filter_)});
+  }
+
+  void SetFilter(std::vector<float> f) { PopulateTensor(filter_, f); }
+
+  void SetBias(std::initializer_list<float> f) { PopulateTensor(bias_, f); }
+
+  void SetInput(std::vector<float> data) { PopulateTensor(input_, data); }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input_;
+  int filter_;
+  int bias_;
+  int output_;
+};
+
+class PrepareOnlyConv3dOpModel : public SingleOpModel {
+ public:
+  PrepareOnlyConv3dOpModel(
+      const TensorData& input, const TensorData& filter,
+      const TensorData& output, Padding padding = Padding_VALID,
+      int32_t stride_depth = 1, int32_t stride_width = 1,
+      int32_t stride_height = 1,
+      ActivationFunctionType activation = ActivationFunctionType_NONE,
+      int32_t dilation_depth = 1, int32_t dilation_width = 1,
+      int32_t dilation_height = 1) {
+    input_ = AddInput(input);
+    filter_ = AddInput(filter);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    resolver_ = std::make_unique<SingleOpResolver>(
+        BuiltinOperator_CONV_3D, ops::builtin::Register_CONV_3D_GENERIC_OPT());
+    BuildInterpreter({GetShape(input_), GetShape(filter_)},
+                     /*num_threads=*/1, /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int filter_;
+  int output_;
+};
+
+template <typename T>
+std::vector<T> CreateRangeVector(int N) {
+  std::vector<T> result;
+  for (int i = 0; i < N; ++i) result.push_back(i);
+  return result;
+}
+
+TEST(Conv3dOpModel, InvalidInputDimsTest) {
+  EXPECT_DEATH_IF_SUPPORTED(Conv3dOpModel m({TensorType_FLOAT32, {2, 2, 4, 1}},
+                                            {TensorType_FLOAT32, {3, 2, 2, 1}},
+                                            {TensorType_FLOAT32, {}}),
+                            "input->dims->size != 5");
+}
+
+TEST(Conv3dOpModel, InvalidFilterDimsTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dOpModel m({TensorType_FLOAT32, {1, 2, 2, 4, 1}},
+                      {TensorType_FLOAT32, {3, 2, 2, 1}},
+                      {TensorType_FLOAT32, {}}),
+      "filter->dims->size != 5");
+}
+
+TEST(Conv3dOpModel, MismatchChannelSizeTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dOpModel m({TensorType_FLOAT32, {1, 2, 2, 4, 1}},
+                      {TensorType_FLOAT32, {1, 3, 2, 2, 2}},
+                      {TensorType_FLOAT32, {}}),
+      "input->dims->data.4. != filter->dims->data.3.");
+}
+
+TEST(Conv3dOpModel, MismatchBiasSizeTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dOpModel m({TensorType_FLOAT32, {1, 2, 2, 4, 2}},
+                      {TensorType_FLOAT32, {1, 3, 2, 2, 1}},
+                      {TensorType_FLOAT32, {2}}, {TensorType_FLOAT32, {}}),
+      "NumElements.bias. != SizeOfDimension.filter, 4.");
+}
+
+TEST(Conv3dPrepareSecurityTest, RejectsIm2ColDepthOverflow) {
+  if (sizeof(void*) <= 4) {
+    GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare "
+                    "on 32-bit.";
+  }
+  constexpr int kHugeDim = 46341;
+  PrepareOnlyConv3dOpModel m(
+      {TensorType_FLOAT32, {1, 1, 1, 1, kHugeDim}},
+      {TensorType_FLOAT32, {kHugeDim, 1, 1, kHugeDim, 1}},
+      {TensorType_FLOAT32, {}}, Padding_SAME);
+
+  // On non-mobile platforms, need_im2col is always true, so the overflow is
+  // detected and the allocation is rejected, resulting in kTfLiteError.
+  // On mobile platforms, it goes to a fallback execution path that does not
+  // require im2col, thus does not overflow and returns kTfLiteOk.
+  if (IsMobilePlatform()) {
+    EXPECT_EQ(m.AllocateTensors(), kTfLiteOk);
+  } else {
+    EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+  }
+}
+
+TEST(Conv3dOpModel, SimpleFloat32Test) {
+  Conv3dOpModel m({TensorType_FLOAT32, {1, 2, 2, 4, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}});
+
+  m.SetInput(CreateRangeVector<float>(32));
+  m.SetFilter({-1, -1, -1, -1, -1, 1, -1, 1, -1, 1,  1,  1, 1, 1,  -1, -1,
+               1,  -1, 1,  1,  1,  1, -1, 1, -1, -1, -1, 1, 1, -1, 1,  -1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 1, 1, 3, 2));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({30, 6, 26, 10, 22, 14}));
+}
+
+TEST(Conv3dOpModel, PaddingValidTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {1, 3, 4, 5, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}});
+
+  m.SetInput(CreateRangeVector<float>(120));
+  m.SetFilter({-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 1,  1, -1, -1,
+               1,  1,  -1, 1,  -1, 1,  -1, 1,  -1, -1, -1, 1, -1, 1, 1,  1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 2, 3, 4, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({-214, 266, -234, 270, -254, 274, -274, 278, -314, 286,
+                        -334, 290, -354, 294, -374, 298, -414, 306, -434, 310,
+                        -454, 314, -474, 318, -614, 346, -634, 350, -654, 354,
+                        -674, 358, -714, 366, -734, 370, -754, 374, -774, 378,
+                        -814, 386, -834, 390, -854, 394, -874, 398}));
+}
+
+TEST(Conv3dOpModel, PaddingSameTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {1, 3, 4, 5, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}}, Padding_SAME);
+
+  m.SetInput(CreateRangeVector<float>(120));
+  m.SetFilter({1,  -1, 1,  -1, 1,  -1, -1, 1, 1, -1, -1, 1, 1,  -1, -1, 1,
+               -1, 1,  -1, 1,  -1, -1, -1, 1, 1, 1,  1,  1, -1, 1,  -1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 3, 4, 5, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-172, 290,  -176, 298,  -180, 306,  -184, 314,  36,   198,  -192,
+           330,  -196, 338,  -200, 346,  -204, 354,  56,   218,  -212, 370,
+           -216, 378,  -220, 386,  -224, 394,  76,   238,  -226, 82,   -230,
+           82,   -234, 82,   -238, 82,   -80,  80,   -252, 450,  -256, 458,
+           -260, 466,  -264, 474,  116,  278,  -272, 490,  -276, 498,  -280,
+           506,  -284, 514,  136,  298,  -292, 530,  -296, 538,  -300, 546,
+           -304, 554,  156,  318,  -306, 82,   -310, 82,   -314, 82,   -318,
+           82,   -80,  80,   158,  -158, 162,  -162, 166,  -166, 170,  -170,
+           176,  -176, 178,  -178, 182,  -182, 186,  -186, 190,  -190, 196,
+           -196, 198,  -198, 202,  -202, 206,  -206, 210,  -210, 216,  -216,
+           220,  -220, 224,  -224, 228,  -228, 232,  -232, 237,  -237}));
+}
+
+TEST(Conv3dOpModel, StrideTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {2, 2, 3, 4, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}}, Padding_VALID, /*stride_depth=*/2,
+                  /*stride_width=*/2, /*stride_height=*/2);
+
+  m.SetInput(CreateRangeVector<float>(96));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 1, 1, 2, 2));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({52, 8, 68, 8, 244, 8, 260, 8}));
+}
+
+TEST(Conv3dOpModel, StrideAndPaddingSameTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {2, 2, 3, 4, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}}, Padding_SAME, /*stride_depth=*/2,
+                  /*stride_width=*/2, /*stride_height=*/2);
+
+  m.SetInput(CreateRangeVector<float>(96));
+  m.SetFilter({-1, 1, -1, 1,  1,  1,  1,  1,  -1, 1, -1, -1, -1, 1,  1,  1,
+               1,  1, -1, -1, -1, -1, -1, -1, 1,  1, 1,  -1, -1, -1, -1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 1, 2, 2, 2));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({-70, -28, -86, -12, -82, -16, -90, -8, -262,
+                                164, -278, 180, -178, 80, -186, 88}));
+}
+
+TEST(Conv3dOpModel, DilationTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {2, 2, 3, 4, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {}}, Padding_VALID, /*stride_depth=*/1,
+                  /*stride_width=*/1, /*stride_height=*/1,
+                  /*activation=*/ActivationFunctionType_NONE,
+                  /*dilation_depth=*/1, /*dilation_width=*/1,
+                  /*dilation_height=*/2);
+
+  m.SetInput(CreateRangeVector<float>(96));
+  m.SetFilter(CreateRangeVector<float>(32));
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 1, 1, 3, 2));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({7248, 7592, 7728, 8104, 8208, 8616, 18768,
+                                19880, 19248, 20392, 19728, 20904}));
+}
+
+TEST(Conv3dOpModel, BiasTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {2, 2, 3, 4, 2}},
+                  {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                  {TensorType_FLOAT32, {2}}, {TensorType_FLOAT32, {}},
+                  Padding_VALID, /*stride_depth=*/2,
+                  /*stride_width=*/2, /*stride_height=*/2);
+
+  m.SetInput(CreateRangeVector<float>(96));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  m.SetBias({1, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 1, 1, 2, 2));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({53, 10, 69, 10, 245, 10, 261, 10}));
+}
+
+TEST(Conv3dOpModel, NoIm2ColTensorTest) {
+  Conv3dOpModel m({TensorType_FLOAT32, {1, 2, 2, 2, 4}},
+                  {TensorType_FLOAT32, {1, 1, 1, 4, 4}},
+                  {TensorType_FLOAT32, {}}, Padding_VALID);
+
+  m.SetInput(CreateRangeVector<float>(32));
+  m.SetFilter(CreateRangeVector<float>(16));
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 2, 2, 2, 4));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray({56,  62,  68,  74,  152, 174, 196, 218, 248, 286, 324,
+                        362, 344, 398, 452, 506, 440, 510, 580, 650, 536, 622,
+                        708, 794, 632, 734, 836, 938, 728, 846, 964, 1082}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/conv3d_transpose.cc

@@ -0,0 +1,348 @@
+/* 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/internal/reference/conv3d_transpose.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/padding.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace conv3d_transpose {
+
+enum KernelType {
+  kReference,
+  kGenericOptimized,
+};
+
+const int kTensorNotAllocated = -1;
+
+struct OpData {
+  Padding3DValues padding;
+
+  // The id of the temporary col2im tensor.
+  int col2im_id = kTensorNotAllocated;
+
+  // The index of col2im tensor in the temporaries list.
+  int col2im_index;
+
+  bool need_col2im = false;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* opdata = new OpData;
+  return opdata;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete static_cast<OpData*>(buffer);
+}
+
+static TfLiteStatus AllocateTemporaryTensorsIfRequired(TfLiteContext* context,
+                                                       TfLiteNode* node,
+                                                       KernelType kernel_type) {
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+  int temporaries_count = 0;
+
+  // Allocate col2im tensor for the optimized kernel.
+  if (kernel_type == kGenericOptimized) {
+    if (data->col2im_id == kTensorNotAllocated) {
+      context->AddTensors(context, 1, &data->col2im_id);
+    }
+    data->col2im_index = temporaries_count++;
+    data->need_col2im = true;
+  }
+
+  TfLiteIntArrayFree(node->temporaries);
+  node->temporaries = TfLiteIntArrayCreate(temporaries_count);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus ResizeOutputAndTemporaryTensors(
+    TfLiteContext* context, OpData* opdata, TfLiteConv3DTransposeParams* params,
+    const TfLiteTensor* shape_tensor, const TfLiteTensor* filter,
+    const TfLiteTensor* input, TfLiteTensor* col2im, TfLiteTensor* output) {
+  auto shape_data = GetTensorData<int32_t>(shape_tensor);
+  for (int i = 0; i < NumElements(shape_tensor); ++i) {
+    TF_LITE_ENSURE_MSG(
+        context, shape_data[i] >= 0, "%s",
+        "Conv3DTranspose output shape has a negative dimension.");
+  }
+  // Output and input tensor must have the same batch size.
+  TF_LITE_ENSURE_EQ(context, shape_data[0], SizeOfDimension(input, 0));
+  // The number of channels of output must be divisible by that of filter.
+  const int filter_output_channels = SizeOfDimension(filter, 3);
+  TF_LITE_ENSURE(context, filter_output_channels != 0);
+  TF_LITE_ENSURE_EQ(context, shape_data[4] % filter_output_channels, 0);
+
+  // Compute padding.
+  const RuntimeShape& filter_shape = GetTensorShape(filter);
+  const int depth = shape_data[1];
+  const int height = shape_data[2];
+  const int width = shape_data[3];
+  const int filter_depth = filter_shape.Dims(0);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  int unused_out_width, unused_out_height, unused_out_depth;
+  opdata->padding = ComputePadding3DValues(
+      params->stride_height, params->stride_width, params->stride_depth,
+      params->dilation_height_factor, params->dilation_width_factor,
+      params->dilation_depth_factor, height, width, depth, filter_height,
+      filter_width, filter_depth, params->padding, &unused_out_height,
+      &unused_out_width, &unused_out_depth);
+  // Computed shape must match the shape of the input tensor.
+  TF_LITE_ENSURE_EQ(context, unused_out_depth, SizeOfDimension(input, 1));
+  TF_LITE_ENSURE_EQ(context, unused_out_height, SizeOfDimension(input, 2));
+  TF_LITE_ENSURE_EQ(context, unused_out_width, SizeOfDimension(input, 3));
+
+  std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> output_shape(
+      TfLiteIntArrayCreate(NumElements(shape_tensor)), TfLiteIntArrayFree);
+  for (int i = 0; i < output_shape->size; ++i) {
+    output_shape->data[i] = GetTensorData<int32_t>(shape_tensor)[i];
+  }
+
+  TF_LITE_ENSURE_STATUS(
+      context->ResizeTensor(context, output, output_shape.release()));
+
+  // Resize col2im tensor.
+  if (opdata->need_col2im) {
+    const RuntimeShape& input_shape = GetTensorShape(input);
+    int col2im_rows = 0;
+    TF_LITE_ENSURE_MSG(context,
+                       input_shape.CheckedNumElementsInRange(
+                           /*start=*/1, /*end=*/4, col2im_rows),
+                       "%s",
+                       "Conv3DTranspose col2im tensor has too many rows.");
+
+    int col2im_columns = 0;
+    TF_LITE_ENSURE_MSG(
+        context, filter_shape.CheckedSizeToDimension(/*end=*/4, col2im_columns),
+        "%s", "Conv3DTranspose col2im tensor has too many columns.");
+
+    std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>
+        col2im_shape_array(TfLiteIntArrayCreate(2), TfLiteIntArrayFree);
+    col2im_shape_array->data[0] = col2im_rows;
+    col2im_shape_array->data[1] = col2im_columns;
+
+    col2im->type = kTfLiteFloat32;
+    col2im->allocation_type = kTfLiteDynamic;
+    return context->ResizeTensor(context, col2im, col2im_shape_array.release());
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
+                     TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteConv3DTransposeParams*>(node->builtin_data);
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+  // Check number of inputs/outputs.
+  TF_LITE_ENSURE(context, node->inputs->size == 3 || node->inputs->size == 4);
+  TF_LITE_ENSURE_EQ(context, node->outputs->size, 1);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  const TfLiteTensor* output_shape;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &output_shape));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 2, &input));
+
+  // Check dimensionality of inputs/outputs.
+  TF_LITE_ENSURE_EQ(context, output_shape->dims->size, 1);
+  TF_LITE_ENSURE_EQ(context, NumElements(output_shape), 5);
+  TF_LITE_ENSURE_EQ(context, input->dims->size, 5);
+  TF_LITE_ENSURE_EQ(context, filter->dims->size, 5);
+
+  // Input and filter must have the same number of channels.
+  TF_LITE_ENSURE_EQ(context, SizeOfDimension(input, 4),
+                    SizeOfDimension(filter, 4));
+
+  // Check types.
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, filter->type, kTfLiteFloat32);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
+  TF_LITE_ENSURE_TYPES_EQ(context, output_shape->type, kTfLiteInt32);
+
+  // Check bias.
+  const TfLiteTensor* bias = GetInput(context, node, 3);
+  if (bias) {
+    TF_LITE_ENSURE_TYPES_EQ(context, bias->type, input->type);
+    TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 3));
+  }
+
+  // GenericOptimized kernel currently doesn't support dilation.
+  if (params->dilation_depth_factor > 1 || params->dilation_height_factor > 1 ||
+      params->dilation_width_factor > 1) {
+    kernel_type = kReference;
+  }
+
+  // Allocate temporary tensors.
+  TF_LITE_ENSURE_STATUS(
+      AllocateTemporaryTensorsIfRequired(context, node, kernel_type));
+
+  // Check temporary tensors.
+  TfLiteTensor* col2im = nullptr;
+  if (opdata->need_col2im) {
+    node->temporaries->data[opdata->col2im_index] = opdata->col2im_id;
+    TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node,
+                                                opdata->col2im_index, &col2im));
+  }
+
+  // Resize the output tensor.
+  if (!IsConstantOrPersistentTensor(output_shape)) {
+    SetTensorToDynamic(output);
+    if (opdata->need_col2im) {
+      SetTensorToDynamic(col2im);
+    }
+  } else {
+    TF_LITE_ENSURE_STATUS(ResizeOutputAndTemporaryTensors(
+        context, opdata, params, output_shape, filter, input, col2im, output));
+  }
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  return Prepare(kernel_type, context, node);
+}
+
+void EvalFloat(KernelType kernel_type, TfLiteContext* context, TfLiteNode* node,
+               TfLiteConv3DTransposeParams* params, OpData* opdata,
+               const TfLiteTensor* input, const TfLiteTensor* filter,
+               const TfLiteTensor* bias, TfLiteTensor* col2im,
+               TfLiteTensor* output) {
+  float output_activation_min, output_activation_max;
+  CalculateActivationRange(params->activation, &output_activation_min,
+                           &output_activation_max);
+
+  Conv3DTransposeParams runtime_params;
+  runtime_params.padding_values = opdata->padding;
+  runtime_params.stride_depth = params->stride_depth;
+  runtime_params.stride_height = params->stride_height;
+  runtime_params.stride_width = params->stride_width;
+  runtime_params.dilation_depth = params->dilation_depth_factor;
+  runtime_params.dilation_height = params->dilation_height_factor;
+  runtime_params.dilation_width = params->dilation_width_factor;
+  runtime_params.float_activation_min = output_activation_min;
+  runtime_params.float_activation_max = output_activation_max;
+
+  switch (kernel_type) {
+    case kReference: {
+      reference_ops::Conv3DTranspose(
+          runtime_params, GetTensorShape(input), GetTensorData<float>(input),
+          GetTensorShape(filter), GetTensorData<float>(filter),
+          GetTensorShape(bias), GetTensorData<float>(bias),
+          GetTensorShape(output), GetTensorData<float>(output));
+      break;
+    }
+    case kGenericOptimized: {
+      optimized_ops::Conv3DTranspose(
+          runtime_params, GetTensorShape(input), GetTensorData<float>(input),
+          GetTensorShape(filter), GetTensorData<float>(filter),
+          GetTensorShape(bias), GetTensorData<float>(bias),
+          GetTensorShape(output), GetTensorData<float>(output),
+          GetTensorShape(col2im), GetTensorData<float>(col2im),
+          CpuBackendContext::GetFromContext(context));
+    } break;
+  }
+}
+
+TfLiteStatus Eval(KernelType kernel_type, TfLiteContext* context,
+                  TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteConv3DTransposeParams*>(node->builtin_data);
+  OpData* opdata = reinterpret_cast<OpData*>(node->user_data);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+  const TfLiteTensor* output_shape;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &output_shape));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 1, &filter));
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 2, &input));
+  const TfLiteTensor* bias = GetInput(context, node, 3);
+  TfLiteTensor* col2im = opdata->need_col2im
+                             ? GetTemporary(context, node, opdata->col2im_index)
+                             : nullptr;
+
+  if (IsDynamicTensor(output)) {
+    TF_LITE_ENSURE_OK(context, ResizeOutputAndTemporaryTensors(
+                                   context, opdata, params, output_shape,
+                                   filter, input, col2im, output));
+  }
+
+  // GenericOptimized kernel currently doesn't support dilation.
+  if (params->dilation_depth_factor > 1 || params->dilation_height_factor > 1 ||
+      params->dilation_width_factor > 1) {
+    kernel_type = kReference;
+  }
+
+  switch (input->type) {
+    case kTfLiteFloat32:
+      EvalFloat(kernel_type, context, node, params, opdata, input, filter, bias,
+                col2im, output);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %s currently not supported.",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  return Eval(kernel_type, context, node);
+}
+
+}  // namespace conv3d_transpose
+
+TfLiteRegistration* Register_CONV_3D_TRANSPOSE_REF() {
+  static TfLiteRegistration r = {
+      conv3d_transpose::Init, conv3d_transpose::Free,
+      conv3d_transpose::Prepare<conv3d_transpose::kReference>,
+      conv3d_transpose::Eval<conv3d_transpose::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONV_3D_TRANSPOSE_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      conv3d_transpose::Init, conv3d_transpose::Free,
+      conv3d_transpose::Prepare<conv3d_transpose::kGenericOptimized>,
+      conv3d_transpose::Eval<conv3d_transpose::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_CONV_3D_TRANSPOSE() {
+  return Register_CONV_3D_TRANSPOSE_GENERIC_OPT();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/conv3d_transpose_test.cc

@@ -0,0 +1,432 @@
+/* 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 <initializer_list>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/kernels/builtin_op_kernels.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+enum class TestType {
+  kConst = 0,
+  kDynamic = 1,
+};
+
+class Conv3dTransposeOpModel : public SingleOpModel {
+ public:
+  Conv3dTransposeOpModel(
+      std::initializer_list<int> output_shape_data, const TensorData& filter,
+      const TensorData& input, const TensorData& bias, const TensorData& output,
+      TestType test_type, Padding padding = Padding_VALID,
+      int32_t stride_depth = 1, int32_t stride_width = 1,
+      int32_t stride_height = 1,
+      ActivationFunctionType activation = ActivationFunctionType_NONE,
+      int32_t dilation_depth = 1, int32_t dilation_width = 1,
+      int32_t dilation_height = 1) {
+    if (test_type == TestType::kDynamic) {
+      output_shape_ = AddInput({TensorType_INT32, {5}});
+    } else {
+      output_shape_ = AddConstInput(TensorType_INT32, output_shape_data, {5});
+    }
+    filter_ = AddInput(filter);
+    input_ = AddInput(input);
+    bias_ = AddInput(bias);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D_TRANSPOSE, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    BuildInterpreter({GetShape(output_shape_), GetShape(filter_),
+                      GetShape(input_), GetShape(bias_)});
+
+    if (test_type == TestType::kDynamic) {
+      PopulateTensor(output_shape_, output_shape_data);
+    }
+  }
+
+  Conv3dTransposeOpModel(
+      std::initializer_list<int> output_shape_data, const TensorData& filter,
+      const TensorData& input, const TensorData& output, TestType test_type,
+      Padding padding = Padding_VALID, int32_t stride_depth = 1,
+      int32_t stride_width = 1, int32_t stride_height = 1,
+      ActivationFunctionType activation = ActivationFunctionType_NONE,
+      int32_t dilation_depth = 1, int32_t dilation_width = 1,
+      int32_t dilation_height = 1) {
+    if (test_type == TestType::kDynamic) {
+      output_shape_ = AddInput({TensorType_INT32, {5}});
+    } else {
+      output_shape_ = AddConstInput(TensorType_INT32, output_shape_data, {5});
+    }
+    filter_ = AddInput(filter);
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D_TRANSPOSE, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    BuildInterpreter(
+        {GetShape(output_shape_), GetShape(filter_), GetShape(input_)});
+    if (test_type == TestType::kDynamic) {
+      PopulateTensor(output_shape_, output_shape_data);
+    }
+  }
+
+  void SetFilter(std::vector<float> f) { PopulateTensor(filter_, f); }
+
+  void SetBias(std::initializer_list<float> f) { PopulateTensor(bias_, f); }
+
+  void SetInput(std::vector<float> data) { PopulateTensor(input_, data); }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int output_shape_;
+  int input_;
+  int filter_;
+  int bias_;
+  int output_;
+};
+
+class PrepareOnlyConv3dTransposeOpModel : public SingleOpModel {
+ public:
+  PrepareOnlyConv3dTransposeOpModel(
+      std::initializer_list<int> output_shape_data, const TensorData& filter,
+      const TensorData& input, const TensorData& output,
+      Padding padding = Padding_VALID, int32_t stride_depth = 1,
+      int32_t stride_width = 1, int32_t stride_height = 1,
+      ActivationFunctionType activation = ActivationFunctionType_NONE,
+      int32_t dilation_depth = 1, int32_t dilation_width = 1,
+      int32_t dilation_height = 1) {
+    output_shape_ = AddConstInput(TensorType_INT32, output_shape_data, {5});
+    filter_ = AddInput(filter);
+    input_ = AddInput(input);
+    output_ = AddOutput(output);
+    SetBuiltinOp(
+        BuiltinOperator_CONV_3D_TRANSPOSE, BuiltinOptions_Conv3DOptions,
+        CreateConv3DOptions(builder_, padding, stride_depth, stride_width,
+                            stride_height, activation, dilation_depth,
+                            dilation_width, dilation_height)
+            .Union());
+    resolver_ = std::make_unique<SingleOpResolver>(
+        BuiltinOperator_CONV_3D_TRANSPOSE,
+        ops::builtin::Register_CONV_3D_TRANSPOSE());
+    BuildInterpreter(
+        {GetShape(output_shape_), GetShape(filter_), GetShape(input_)},
+        /*num_threads=*/1, /*allow_fp32_relax_to_fp16=*/false,
+        /*apply_delegate=*/false,
+        /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int output_shape_;
+  int input_;
+  int filter_;
+  int output_;
+};
+
+template <typename T>
+std::vector<T> CreateRangeVector(int N) {
+  std::vector<T> result;
+  for (int i = 0; i < N; ++i) result.push_back(i);
+  return result;
+}
+
+class Conv3dTransposeOpTest : public ::testing::TestWithParam<TestType> {};
+
+TEST_P(Conv3dTransposeOpTest, InvalidInputDimsTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dTransposeOpModel m(
+          {1, 2, 3, 4, 5}, {TensorType_FLOAT32, {2, 2, 4, 1}},
+          {TensorType_FLOAT32, {3, 2, 2, 1}}, {TensorType_FLOAT32, {}},
+          Conv3dTransposeOpTest::GetParam()),
+      "input->dims->size != 5");
+}
+
+TEST_P(Conv3dTransposeOpTest, InvalidFilterDimsTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dTransposeOpModel m(
+          {1, 2, 3, 4, 5}, {TensorType_FLOAT32, {2, 2, 4, 1}},
+          {TensorType_FLOAT32, {1, 3, 2, 2, 1}}, {TensorType_FLOAT32, {}},
+          Conv3dTransposeOpTest::GetParam()),
+      "filter->dims->size != 5");
+}
+
+TEST_P(Conv3dTransposeOpTest, MismatchChannelSizeTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dTransposeOpModel m(
+          {1, 2, 3, 4, 5}, {TensorType_FLOAT32, {1, 2, 2, 4, 1}},
+          {TensorType_FLOAT32, {1, 3, 2, 2, 2}}, {TensorType_FLOAT32, {}},
+          Conv3dTransposeOpTest::GetParam()),
+      "SizeOfDimension.input, 4. != SizeOfDimension.filter, 4.");
+}
+
+TEST_P(Conv3dTransposeOpTest, MismatchBiasSizeTest) {
+  EXPECT_DEATH_IF_SUPPORTED(
+      Conv3dTransposeOpModel m(
+          {1, 2, 3, 4, 5}, {TensorType_FLOAT32, {1, 3, 2, 2, 2}},
+          {TensorType_FLOAT32, {1, 2, 2, 4, 2}}, {TensorType_FLOAT32, {3}},
+          {TensorType_FLOAT32, {}}, Conv3dTransposeOpTest::GetParam()),
+      "NumElements.bias. != SizeOfDimension.filter, 3.");
+}
+
+TEST(Conv3dTransposePrepareSecurityTest, RejectsCol2ImOverflow) {
+  if (sizeof(void*) <= 4) {
+    GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare "
+                    "on 32-bit.";
+  }
+  constexpr int kHugeDim = 46341;
+  PrepareOnlyConv3dTransposeOpModel m(
+      {1, 1, 1, 1, 0}, {TensorType_FLOAT32, {1, kHugeDim, kHugeDim, 1, 0}},
+      {TensorType_FLOAT32, {1, 1, 1, 1, 0}}, {TensorType_FLOAT32, {}},
+      Padding_SAME);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
+TEST(Conv3dTransposePrepareSecurityTest, RejectsZeroFilterOutputChannels) {
+  PrepareOnlyConv3dTransposeOpModel m({1, 1, 1, 1, 1},
+                                      {TensorType_FLOAT32, {1, 1, 1, 0, 1}},
+                                      {TensorType_FLOAT32, {1, 1, 1, 1, 1}},
+                                      {TensorType_FLOAT32, {}}, Padding_SAME);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
+TEST_P(Conv3dTransposeOpTest, SimpleFloat32Test) {
+  Conv3dTransposeOpModel m(
+      {1, 3, 3, 5, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {1, 2, 2, 4, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam());
+
+  m.SetInput(CreateRangeVector<float>(32));
+  m.SetFilter({-1, -1, -1, -1, -1, 1, -1, 1, -1, 1,  1,  1, 1, 1,  -1, -1,
+               1,  -1, 1,  1,  1,  1, -1, 1, -1, -1, -1, 1, 1, -1, 1,  -1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 3, 3, 5, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-1,  -1,  -4,  -4,  -8,  -8,  -12, -12, 1,   1,   -16, -16, -18,
+           -16, -18, -20, -18, -24, 14,  -12, 1,   17,  18,  4,   22,  4,
+           26,  4,   29,  -29, -34, -32, -36, -30, -36, -30, -36, -30, 14,
+           2,   -50, 2,   -8,  -26, -8,  -26, -8,  -26, 74,  -44, -16, 50,
+           28,  4,   28,  4,   28,  4,   60,  -62, -1,  33,  32,  38,  36,
+           42,  40,  46,  45,  1,   -34, 50,  10,  54,  10,  58,  10,  62,
+           60,  0,   -49, 1,   -54, 0,   -58, 0,   -62, 0,   -1,  -1}));
+}
+
+TEST_P(Conv3dTransposeOpTest, PaddingValidTest) {
+  Conv3dTransposeOpModel m(
+      {1, 4, 5, 6, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {1, 3, 4, 5, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam());
+
+  m.SetInput(CreateRangeVector<float>(120));
+  m.SetFilter({-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 1, 1,  1, -1, -1,
+               1,  1,  -1, 1,  -1, 1,  -1, 1,  -1, -1, -1, 1, -1, 1, 1,  1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 4, 5, 6, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-1,   -1,   -6,   -6,   -14,  -14,  -22,  -22,  -30,  -30,  -17,
+           -17,  -22,  -20,  -50,  -46,  -58,  -58,  -66,  -70,  -74,  -82,
+           -20,  -54,  -62,  -40,  -90,  -106, -98,  -118, -106, -130, -114,
+           -142, -20,  -94,  -102, -60,  -130, -166, -138, -178, -146, -190,
+           -154, -202, -20,  -134, -61,  1,    -4,   -60,  -4,   -64,  -4,
+           -68,  -4,   -72,  77,   -77,  -80,  -80,  -160, -164, -164, -172,
+           -168, -180, -172, -188, -96,  -96,  -162, -98,  -188, -282, -196,
+           -290, -204, -298, -212, -306, -18,  -196, -202, -118, -228, -322,
+           -236, -330, -244, -338, -252, -346, -18,  -216, -242, -138, -268,
+           -362, -276, -370, -284, -378, -292, -386, -18,  -236, -202, 2,
+           -68,  -78,  -72,  -78,  -76,  -78,  -80,  -78,  158,  -80,  -80,
+           -160, -240, -324, -244, -332, -248, -340, -252, -348, -176, -176,
+           -322, -178, -348, -442, -356, -450, -364, -458, -372, -466, -18,
+           -276, -362, -198, -388, -482, -396, -490, -404, -498, -412, -506,
+           -18,  -296, -402, -218, -428, -522, -436, -530, -444, -538, -452,
+           -546, -18,  -316, -362, 2,    -148, -78,  -152, -78,  -156, -78,
+           -160, -78,  238,  -80,  161,  1,    166,  2,    170,  2,    174,
+           2,    178,  2,    1,    1,    20,   2,    22,   164,  22,   168,
+           22,   172,  22,   176,  2,    178,  20,   2,    22,   184,  22,
+           188,  22,   192,  22,   196,  2,    198,  20,   2,    22,   204,
+           22,   208,  22,   212,  22,   216,  2,    218,  -221, 1,    -224,
+           222,  -228, 226,  -232, 230,  -236, 234,  1,    237}));
+}
+
+TEST_P(Conv3dTransposeOpTest, PaddingSameTest) {
+  Conv3dTransposeOpModel m(
+      {1, 3, 4, 5, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {1, 3, 4, 5, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam(), Padding_SAME);
+
+  m.SetInput(CreateRangeVector<float>(120));
+  m.SetFilter({1,  -1, 1,  -1, 1,  -1, -1, 1, 1, -1, -1, 1, 1,  -1, -1, 1,
+               -1, 1,  -1, 1,  -1, -1, -1, 1, 1, 1,  1,  1, -1, 1,  -1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 3, 4, 5, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-1,  -1,  -2,  0,   -2,  0,   -2,  0,   -2,  0,   -2,  0,   -4,  2,
+           -4,  2,   -4,  2,   -4,  2,   -2,  0,   -4,  2,   -4,  2,   -4,  2,
+           -4,  2,   -2,  0,   -4,  2,   -4,  2,   -4,  2,   -4,  2,   0,   0,
+           -2,  2,   -6,  2,   -10, 2,   -14, 2,   0,   2,   -18, 10,  -18, 14,
+           -18, 18,  -18, 22,  20,  22,  -18, 30,  -18, 34,  -18, 38,  -18, 42,
+           40,  42,  -18, 50,  -18, 54,  -18, 58,  -18, 62,  0,   0,   -82, 2,
+           -86, 2,   -90, 2,   -94, 2,   80,  82,  -18, 90,  -18, 94,  -18, 98,
+           -18, 102, 100, 102, -18, 110, -18, 114, -18, 118, -18, 122, 120, 122,
+           -18, 130, -18, 134, -18, 138, -18, 142}));
+}
+
+TEST_P(Conv3dTransposeOpTest, PaddingValidComplexTest) {
+  Conv3dTransposeOpModel m(
+      {2, 4, 3, 2, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {2, 3, 2, 1, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam(), Padding_VALID);
+
+  m.SetInput(CreateRangeVector<float>(24));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 4, 3, 2, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-1, 1,   1, -1, -2, 4,   2, 0,  -1, -5,  1, 5,  -2, 10,  2, -2,
+           -4, 8,   4, 8,  -2, -18, 2, 18, -2, 26,  2, -2, -4, 8,   4, 24,
+           -2, -34, 2, 34, -1, 17,  1, -1, -2, 4,   2, 16, -1, -21, 1, 21,
+           -1, 25,  1, -1, -2, 4,   2, 24, -1, -29, 1, 29, -2, 58,  2, -2,
+           -4, 8,   4, 56, -2, -66, 2, 66, -2, 74,  2, -2, -4, 8,   4, 72,
+           -2, -82, 2, 82, -1, 41,  1, -1, -2, 4,   2, 40, -1, -45, 1, 45}));
+}
+
+TEST_P(Conv3dTransposeOpTest, StrideTest) {
+  Conv3dTransposeOpModel m(
+      {2, 4, 3, 2, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {2, 2, 2, 1, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam(), Padding_VALID,
+      /*stride_depth=*/2,
+      /*stride_width=*/1, /*stride_height=*/1);
+
+  m.SetInput(CreateRangeVector<float>(16));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 4, 3, 2, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {-1, 1,   1, -1, -2, 4,   2, 0,  -1, -5,  1, 5,  -1, 1,   1, -1,
+           -2, 4,   2, 0,  -1, -5,  1, 5,  -1, 9,   1, -1, -2, 4,   2, 8,
+           -1, -13, 1, 13, -1, 9,   1, -1, -2, 4,   2, 8,  -1, -13, 1, 13,
+           -1, 17,  1, -1, -2, 4,   2, 16, -1, -21, 1, 21, -1, 17,  1, -1,
+           -2, 4,   2, 16, -1, -21, 1, 21, -1, 25,  1, -1, -2, 4,   2, 24,
+           -1, -29, 1, 29, -1, 25,  1, -1, -2, 4,   2, 24, -1, -29, 1, 29}));
+}
+
+TEST_P(Conv3dTransposeOpTest, StrideAndPaddingSameTest) {
+  Conv3dTransposeOpModel m(
+      {2, 4, 2, 1, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+      {TensorType_FLOAT32, {2, 2, 2, 1, 2}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam(), Padding_SAME,
+      /*stride_depth=*/2,
+      /*stride_width=*/1, /*stride_height=*/1);
+
+  m.SetInput(CreateRangeVector<float>(16));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 4, 2, 1, 2));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({-1, 1,  -2, 4,  -1, 1,  -2, 4,  -1, 9,  -2,
+                                4,  -1, 9,  -2, 4,  -1, 17, -2, 4,  -1, 17,
+                                -2, 4,  -1, 25, -2, 4,  -1, 25, -2, 4}));
+}
+
+TEST_P(Conv3dTransposeOpTest, DilationTest) {
+  Conv3dTransposeOpModel m(
+      {1, 3, 3, 2, 2}, {TensorType_FLOAT32, {1, 2, 2, 2, 1}},
+      {TensorType_FLOAT32, {1, 3, 1, 1, 1}}, {TensorType_FLOAT32, {}},
+      Conv3dTransposeOpTest::GetParam(), Padding_VALID,
+      /*stride_depth=*/1,
+      /*stride_width=*/1, /*stride_height=*/1,
+      /*activation=*/ActivationFunctionType_NONE,
+      /*dilation_depth=*/1, /*dilation_width=*/1,
+      /*dilation_height=*/2);
+
+  m.SetInput(CreateRangeVector<float>(3));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 3, 3, 2, 2));
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray({0, 0,  0, 0, 0, 0, 0, 0, 0,  0, 0, 0,
+                                1, -1, 1, 1, 0, 0, 0, 0, -1, 1, 1, -1,
+                                2, -2, 2, 2, 0, 0, 0, 0, -2, 2, 2, -2}));
+}
+
+TEST_P(Conv3dTransposeOpTest, BiasTest) {
+  Conv3dTransposeOpModel m({2, 4, 3, 2, 2},
+                           {TensorType_FLOAT32, {2, 2, 2, 2, 2}},
+                           {TensorType_FLOAT32, {2, 3, 2, 1, 2}},
+                           {TensorType_FLOAT32, {2}}, {TensorType_FLOAT32, {}},
+                           Conv3dTransposeOpTest::GetParam(), Padding_VALID);
+
+  m.SetInput(CreateRangeVector<float>(24));
+  m.SetFilter({1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1,
+               1, -1, 1, 1, -1, 1, 1, -1, 1, -1, -1, -1, -1, 1, 1, 1});
+  m.SetBias({1, 2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(2, 4, 3, 2, 2));
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(
+          {0,  3,   2, 1,  -1, 6,   3, 2,  0,  -3,  2, 7,  -1, 12,  3, 0,
+           -3, 10,  5, 10, -1, -16, 3, 20, -1, 28,  3, 0,  -3, 10,  5, 26,
+           -1, -32, 3, 36, 0,  19,  2, 1,  -1, 6,   3, 18, 0,  -19, 2, 23,
+           0,  27,  2, 1,  -1, 6,   3, 26, 0,  -27, 2, 31, -1, 60,  3, 0,
+           -3, 10,  5, 58, -1, -64, 3, 68, -1, 76,  3, 0,  -3, 10,  5, 74,
+           -1, -80, 3, 84, 0,  43,  2, 1,  -1, 6,   3, 42, 0,  -43, 2, 47}));
+}
+
+INSTANTIATE_TEST_SUITE_P(Conv3dTransposeOpTest, Conv3dTransposeOpTest,
+                         ::testing::Values(TestType::kConst,
+                                           TestType::kDynamic));
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/conv_mem_test.cc

@@ -0,0 +1,81 @@
+/* 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 <string>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/macros.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/tools/logging.h"
+
+namespace tflite {
+
+void TestMemoryThreshold(const std::string& model_path,
+                         size_t threshold_in_kb) {
+  // The Im2Col optimization is only applied on mobile platforms, so only
+  // validate on such platforms.
+  if (!IsMobilePlatform()) {
+    return;
+  }
+
+  // The model has a conv op will require a huge temporary tensor if
+  // im2col is performed and it's possible to cause OOM on devices. To prevent
+  // this from happening, a size cap (i.e. kMaxIm2colBufferSizeMobile) of
+  // to-be-allocated im2col data is used to determine whether to disable
+  // im2col. This test will check the memory footprint before/after
+  // interpreter Invoke to ensure the size cap is correctly enforced on mobile
+  // platforms.
+  auto model = FlatBufferModel::BuildFromFile(model_path.c_str());
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+
+  // Note that we explicitly set 1 thread here to avoid extra memory footprint
+  // caused by multithreading, which will make the memory usage threshold
+  // check later more reliable.
+  ASSERT_EQ(InterpreterBuilder(*model, ops::builtin::BuiltinOpResolver())(
+                &interpreter, /*num_threads*/ 1),
+            kTfLiteOk);
+  ASSERT_TRUE(interpreter);
+  ASSERT_EQ(interpreter->AllocateTensors(), kTfLiteOk);
+
+  // Memory required for all tensors should be smaller than the  threshold.
+  int64_t accumulate_tensor_memory = 0;
+  for (int i = 0; i < interpreter->tensors_size(); ++i) {
+    accumulate_tensor_memory += interpreter->tensor(i)->bytes;
+  }
+  EXPECT_LE(accumulate_tensor_memory, threshold_in_kb * 1024);
+}
+
+TEST(ConvMemUsage, HugeIm2ColData) {
+  TestMemoryThreshold(
+      // The model has a conv op will require a temporary tensor of ~3.5GB if
+      // im2col is performed.
+      "tensorflow/lite/testdata/conv_huge_im2col.bin",
+      /*threshold_in_kb=*/3 * 1024 * 1024);
+}
+
+TEST(Conv3DMemUsage, HugeIm2ColData) {
+  TestMemoryThreshold(
+      // The model has a Conv3D op will require a temporary tensor of ~1.3GB if
+      // im2col is performed.If not, it will use about 450MB.
+      "tensorflow/lite/testdata/conv3d_huge_im2col.bin",
+      /*threshold_in_kb=*/1 * 1024 * 1024);
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/cpu_backend_context.cc

@@ -0,0 +1,184 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+
+#include <memory>
+
+#ifdef TFLITE_KERNEL_USE_XNNPACK
+#include "pthreadpool.h"  // from @pthreadpool
+#endif
+
+#ifdef TFLITE_HAVE_CPUINFO
+#include "include/cpuinfo.h"
+#endif
+
+#include "public/gemmlowp.h"
+#include "ruy/context.h"  // from @ruy
+#include "ruy/path.h"  // from @ruy
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/macros.h"
+#include "tensorflow/lite/external_cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+
+namespace {
+const int kDefaultNumThreadpoolThreads = 1;
+
+}  // namespace
+
+namespace tflite {
+
+// Use weak symbols if possible to dispatch to deprecated paths.
+#if TFLITE_HAS_ATTRIBUTE_WEAK && !defined(__APPLE__)
+extern TFLITE_ATTRIBUTE_WEAK bool UseGemmlowpOnX86();
+#endif  // defined(TFLITE_HAS_ATTRIBUTE_WEAK) && !(__APPLE__)
+
+#if defined(TFLITE_HAVE_CPUINFO)
+CpuBackendContext::CpuInfo::~CpuInfo() {
+  if (init_status_ == InitStatus::kInitialized) {
+    cpuinfo_deinitialize();
+  }
+}
+
+bool CpuBackendContext::CpuInfo::EnsureInitialized() {
+  if (init_status_ == InitStatus::kNotYetAttempted) {
+    init_status_ = Initialize();
+  }
+  return init_status_ == InitStatus::kInitialized;
+}
+
+CpuBackendContext::CpuInfo::InitStatus
+CpuBackendContext::CpuInfo::Initialize() {
+  TFLITE_DCHECK_EQ(init_status_, InitStatus::kNotYetAttempted);
+  if (!cpuinfo_initialize()) {
+    return InitStatus::kFailed;
+  }
+  return InitStatus::kInitialized;
+}
+
+bool CpuBackendContext::CpuInfo::Avx2Fma() {
+  return EnsureInitialized() && cpuinfo_has_x86_avx2() &&
+         cpuinfo_has_x86_fma3();
+}
+
+bool CpuBackendContext::CpuInfo::Avx() {
+  return EnsureInitialized() && cpuinfo_has_x86_avx();
+}
+
+bool CpuBackendContext::CpuInfo::Avx512() {
+  return EnsureInitialized() && cpuinfo_has_x86_avx512f() &&
+         cpuinfo_has_x86_avx512dq() && cpuinfo_has_x86_avx512cd() &&
+         cpuinfo_has_x86_avx512bw() && cpuinfo_has_x86_avx512vl();
+}
+#else
+
+CpuBackendContext::CpuInfo::~CpuInfo() {}
+
+bool CpuBackendContext::CpuInfo::EnsureInitialized() {
+  if (init_status_ == InitStatus::kNotYetAttempted) {
+    init_status_ = InitStatus::kInitialized;
+  }
+  TFLITE_DCHECK_EQ(init_status_, InitStatus::kInitialized);
+  return true;
+}
+
+bool CpuBackendContext::CpuInfo::Avx2Fma() { return false; }
+
+bool CpuBackendContext::CpuInfo::Avx() { return false; }
+
+bool CpuBackendContext::CpuInfo::Avx512() { return false; }
+#endif  // TFLITE_HAVE_CPUINFO
+
+CpuBackendContext* CpuBackendContext::GetFromContext(TfLiteContext* context) {
+  auto* external_context = static_cast<ExternalCpuBackendContext*>(
+      context->GetExternalContext(context, kTfLiteCpuBackendContext));
+
+  if (external_context == nullptr) {
+    TF_LITE_FATAL(
+        "ExternalCpuBackendContext isn't properly initialized during TFLite "
+        "interpreter initialization.");
+  }
+
+  auto* cpu_backend_context = static_cast<CpuBackendContext*>(
+      external_context->internal_backend_context());
+  if (cpu_backend_context == nullptr) {
+    // We do the lazy initialization here for the TfLiteInternalBackendContext
+    // that's wrapped inside ExternalCpuBackendContext.
+    cpu_backend_context = new CpuBackendContext();
+    cpu_backend_context->SetMaxNumThreads(context->recommended_num_threads);
+    external_context->set_internal_backend_context(
+        std::unique_ptr<TfLiteInternalBackendContext>(cpu_backend_context));
+  }
+
+  return cpu_backend_context;
+}
+
+CpuBackendContext::CpuBackendContext()
+    : TfLiteInternalBackendContext(),
+      ruy_context_(new ruy::Context),
+      gemmlowp_context_(new gemmlowp::GemmContext) {
+  SetMaxNumThreads(kDefaultNumThreadpoolThreads);
+// TODO(b/148289189) Remove when clients have transitioned to runtime flag.
+#ifdef TFLITE_WITH_RUY_GEMV
+  SetUseCaching(true);
+#else
+  SetUseCaching(false);
+#endif
+}
+
+CpuBackendContext::~CpuBackendContext() {}
+
+void CpuBackendContext::SetMaxNumThreads(int max_num_threads) {
+  const int target_num_threads =
+      max_num_threads > -1 ? max_num_threads : kDefaultNumThreadpoolThreads;
+  max_num_threads_ = target_num_threads;
+  ruy_context_->set_max_num_threads(target_num_threads);
+  gemmlowp_context_->set_max_num_threads(target_num_threads);
+}
+
+void CpuBackendContext::SetUseCaching(bool flag) { use_caching_ = flag; }
+
+#ifdef TFLITE_KERNEL_USE_XNNPACK
+pthreadpool_t CpuBackendContext::get_xnnpack_threadpool() {
+  if (!xnnpack_threadpool_ && max_num_threads_ > 1) {
+    xnnpack_threadpool_.reset(
+        pthreadpool_create(static_cast<size_t>(max_num_threads_)));
+  }
+  return xnnpack_threadpool_.get();
+}
+#endif
+
+bool CpuBackendContext::PreferGemmlowpOnX86() {
+  bool use_gemmlowp_on_x86 = false;
+#if defined(TFLITE_X86_PLATFORM) && TFLITE_HAS_ATTRIBUTE_WEAK && \
+    !defined(__APPLE__)
+  if (::tflite::UseGemmlowpOnX86 != nullptr) {
+    use_gemmlowp_on_x86 = ::tflite::UseGemmlowpOnX86();
+  }
+#endif  // TFLITE_X86_PLATFORM && TFLITE_HAS_ATTRIBUTE_WEAK && !(__APPLE__)
+  return use_gemmlowp_on_x86 || !RuyHasAvxOrAbove();
+}
+
+bool CpuBackendContext::RuyHasAvxOrAbove() {
+  // TODO(b/183178387): Use a proper query to detect AVX/optimized paths.
+#if RUY_PLATFORM_X86_ENHANCEMENTS
+  return cpuinfo_.Avx() || cpuinfo_.Avx2Fma() || cpuinfo_.Avx512();
+#else
+  return false;
+#endif
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/cpu_backend_context.h

@@ -0,0 +1,137 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_CONTEXT_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_CONTEXT_H_
+
+#if (defined(__i386) || defined(_M_IX86) || defined(__x86_64__) || \
+     defined(_M_X64))
+#define TFLITE_X86_PLATFORM
+#endif
+
+#include <memory>
+
+#include "public/gemmlowp.h"
+#ifdef TFLITE_KERNEL_USE_XNNPACK
+#include "pthreadpool.h"  // from @pthreadpool
+#endif
+#include "ruy/context.h"  // from @ruy
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/external_cpu_backend_context.h"
+
+namespace tflite {
+
+class CpuBackendContext final : public TfLiteInternalBackendContext {
+ public:
+  static CpuBackendContext* GetFromContext(TfLiteContext* context);
+
+  CpuBackendContext();
+  ~CpuBackendContext() override;
+
+  ruy::Context* ruy_context() const { return ruy_context_.get(); }
+
+  gemmlowp::GemmContext* gemmlowp_context() const {
+    return gemmlowp_context_.get();
+  }
+
+  // Sets the maximum-number-of-threads-to-use parameter, only as a means of
+  // passing around this information.
+  void SetMaxNumThreads(int max_num_threads) override;
+
+  int max_num_threads() const { return max_num_threads_; }
+
+  void SetUseCaching(bool flag);
+
+  bool use_caching() const { return use_caching_; }
+
+#ifdef TFLITE_KERNEL_USE_XNNPACK
+  pthreadpool_t get_xnnpack_threadpool();
+#endif
+
+  void ClearCaches() override { ruy_context_->ClearPrepackedCache(); }
+
+  // Gemmlowp on x86 is a deprecated path but some clients may still use
+  // this path based on link time dependencies.
+  bool PreferGemmlowpOnX86();
+
+ private:
+  bool RuyHasAvxOrAbove();
+
+  // Copy the wrapper class for cpuinfo from Ruy.
+  class CpuInfo final {
+   public:
+    CpuInfo() {}
+    ~CpuInfo();
+
+    // X86 features
+    bool Avx();
+    bool Avx2Fma();
+    bool Avx512();
+
+   private:
+    enum class InitStatus {
+      kNotYetAttempted,
+      kInitialized,
+      kFailed,
+    };
+
+    InitStatus init_status_ = InitStatus::kNotYetAttempted;
+
+    bool EnsureInitialized();
+    InitStatus Initialize();
+    CpuInfo(const CpuInfo&) = delete;
+    CpuInfo& operator=(const CpuInfo&) = delete;
+  };
+
+  // To enable a smooth transition from the current direct usage
+  // of the underlying gemmlowp context to going through abstractions
+  // (see :cpu_backend_gemm), for now a CpuBackendContext always
+  // stores both a gemmlowp context and a ruy context.
+  // TODO(b/131416458): Once call sites all go through abstractions,
+  // elide what can be elided based on TFLITE_WITH_RUY.
+  const std::unique_ptr<ruy::Context> ruy_context_;
+  const std::unique_ptr<gemmlowp::GemmContext> gemmlowp_context_;
+  CpuInfo cpuinfo_;
+
+  // The maximum of threads used for parallelizing TfLite ops. However,
+  // cpu_backend_threadpool::Execute creates as many threads as it's
+  // asked to, regardless of this. Typically a call site would query
+  // cpu_backend_context->max_num_threads() and used that to determine
+  // the number of tasks to create and to give to
+  // cpu_backend_threadpool::Execute.
+  //
+  // This value also gets propagated to back-ends, where it plays the same
+  // information-only role.
+  int max_num_threads_;
+  // For matrix muliplications with constants parameters (i.e. weights), we can
+  // sometimes provide speedups by caching the "prepacked" data, for some
+  // additional memory cost. This flag permits the user to route all
+  // CpuBackendGem operations to a library that permits such an optimization
+  // (currently the Ruy library only).
+  bool use_caching_;
+
+#ifdef TFLITE_KERNEL_USE_XNNPACK
+  // A smart pointer for the xnnpack threadpool. Is created by a call from the
+  // interpreter, and then consumed by xnnpack, possibly via a TFLite kernel.
+  std::unique_ptr<pthreadpool, decltype(&pthreadpool_destroy)>
+      xnnpack_threadpool_{nullptr, &pthreadpool_destroy};
+#endif
+
+  CpuBackendContext(const CpuBackendContext&) = delete;
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_CONTEXT_H_

tensorflow/lite/kernels/cpu_backend_gemm.h

@@ -0,0 +1,226 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_
+
+#include <cstdint>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
+
+#ifndef TFLITE_WITH_RUY
+#include "tensorflow/lite/kernels/cpu_backend_gemm_eigen.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_x86.h"
+#endif
+
+namespace tflite {
+
+namespace cpu_backend_gemm {
+
+// The main entry point for CpuBackendGemm::Gemm.
+//
+// If TFLITE_WITH_RUY is set, CpuBackendGemm::Gemm will always go to Ruy aka
+// GemmImplUsingRuy. The behavior is as follows:
+//
+//                    |Quantized (uint8)|Quantized (int8)| Float |
+// TFLITE_WITH_RUY    |      Ruy        |      Ruy       | Ruy   |
+// !TFLITE_WITH_RUY   |      gemmlowp   |  Ruy/gemmlowp* | eigen |
+// * - Ruy if NEON is not available.
+//
+//  On most ARM32/ARM64 platforms, the default is TFLITE_WITH_RUY:
+//  (default)         |      Ruy        |     Ruy        | Ruy   |
+//
+//  On other platforms (including x86), the default is !TFLITE_WITH_RUY:
+//  (default)         |      gemmlowp   |     Ruy        | eigen |
+//
+// Use --define=tflite_with_ruy=true or --define=tflite_with_ruy=false to
+// override the default.
+
+#if !defined(TFLITE_WITH_RUY) && defined(TFLITE_X86_PLATFORM)
+/* GEMM dispatch implementation for x86.
+ */
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImpl : detail::GemmImplX86<LhsScalar, RhsScalar, AccumScalar,
+                                      DstScalar, quantization_flavor> {};
+#else
+/* Generic implementation using ruy.
+ * Non-ruy implementation will be partial specializations of this template.
+ */
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImpl : detail::GemmImplUsingRuy<LhsScalar, RhsScalar, AccumScalar,
+                                           DstScalar, quantization_flavor> {};
+
+#if !defined(TFLITE_WITH_RUY)
+
+/* Specializations using gemmlowp */
+template <typename SrcScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+struct GemmImpl<SrcScalar, SrcScalar, std::int32_t, DstScalar,
+                quantization_flavor>
+    : detail::GemmImplUsingGemmlowp<SrcScalar, SrcScalar, std::int32_t,
+                                    DstScalar, quantization_flavor> {};
+
+// When SrcScalar=int8 or DstScalar=int8, gemmlowp fails to compile
+// outside of NEON. We avoid the compilation failure by subspecializing these
+// cases, rerouting it back to ruy.
+#if !defined(GEMMLOWP_NEON)
+template <typename SrcScalar, QuantizationFlavor quantization_flavor>
+struct GemmImpl<SrcScalar, SrcScalar, std::int32_t, std::int8_t,
+                quantization_flavor>
+    : detail::GemmImplUsingRuy<SrcScalar, SrcScalar, std::int32_t, std::int8_t,
+                               quantization_flavor> {};
+
+template <typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImpl<std::int8_t, std::int8_t, std::int32_t, DstScalar,
+                quantization_flavor>
+    : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t,
+                               DstScalar, quantization_flavor> {};
+
+template <QuantizationFlavor quantization_flavor>
+struct GemmImpl<std::int8_t, std::int8_t, std::int32_t, std::int8_t,
+                quantization_flavor>
+    : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t,
+                               std::int8_t, quantization_flavor> {};
+#endif  // not GEMMLOWP_NEON
+
+/* Specializations using Eigen */
+
+template <>
+struct GemmImpl<float, float, float, float, QuantizationFlavor::kFloatingPoint>
+    : detail::GemmImplUsingEigen {};
+
+#endif  // not TFLITE_WITH_RUY
+
+#endif  // not TFLITE_WITH_RUY and TFLITE_X86_PLATFORM
+
+/* Public entry point */
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+void Gemm(const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+          const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+          const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+          const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+          CpuBackendContext* context) {
+  ruy::profiler::ScopeLabel label("cpu_backend_gemm::Gemm");
+  ValidateParams(lhs_params, rhs_params, dst_params, params);
+  if (!IsValidGemm(lhs_params, rhs_params, dst_params)) {
+    // For now, assert in debug mode, return in opt.
+    // TODO(b/183099395) Eliminate debug/release discrepancy by plumbing in
+    // TFLiteStatus so we can return an error here.
+    TFLITE_DCHECK(false);
+    return;
+  }
+  // In some cases we want to unconditionally use ruy as the backend, overriding
+  // the `tflite_with_ruy` setting and the platform default.
+  bool must_use_ruy = false;
+  if (context->use_caching()) {
+    // Only ruy supports caching of pre-packed matrices. Due to the large
+    // performance impact in the cases where it's typically used, this overrides
+    // the default.
+    must_use_ruy = true;
+  }
+  if (lhs_params.order != Order::kRowMajor ||
+      rhs_params.order != Order::kColMajor ||
+      dst_params.order != Order::kColMajor) {
+    // ruy supports all 2^3=8 combinations of storage orders with comparable
+    // performance. In ruy, it's only a runtime switch. In other backends
+    // (gemmlowp, Eigen), storage orders are template parameters, supporting
+    // all 8 combinations would be up to a 8-fold code size increase, so we
+    // prefer to force usage of ruy in these cases.
+    must_use_ruy = true;
+  }
+  if (must_use_ruy) {
+    detail::GemmImplUsingRuy<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                             quantization_flavor>::Run(lhs_params, lhs_data,
+                                                       rhs_params, rhs_data,
+                                                       dst_params, dst_data,
+                                                       params, context);
+    return;
+  }
+  // If we did not choose to force usage of ruy above, then we may now consider
+  // using custom GEMV code for the matrix*vector cases.
+  const bool try_custom_gemv = (dst_params.cols == 1);
+  if (try_custom_gemv) {
+    // GEMV case: try a custom fast GEMV path. It will return true if it
+    // actually handled it.
+    if (detail::CustomGemv(lhs_params, lhs_data, rhs_params, rhs_data,
+                           dst_params, dst_data, params, context)) {
+      return;
+    }
+  }
+  // Generic case: dispatch to any backend as a general GEMM.
+  GemmImpl<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+           quantization_flavor>::Run(lhs_params, lhs_data, rhs_params, rhs_data,
+                                     dst_params, dst_data, params, context);
+}
+
+// Special path for 16x8 quant gemm.
+template <QuantizationFlavor quantization_flavor>
+void Gemm(const MatrixParams<int8_t>& lhs_params, const int8_t* lhs_data,
+          const MatrixParams<int16_t>& rhs_params, const int16_t* rhs_data,
+          const MatrixParams<int16_t>& dst_params, int16_t* dst_data,
+          const GemmParams<int32_t, int16_t, quantization_flavor>& params,
+          CpuBackendContext* context) {
+  ruy::profiler::ScopeLabel label("cpu_backend_gemm::Gemm");
+  ValidateParams(lhs_params, rhs_params, dst_params, params);
+  if (!IsValidGemm(lhs_params, rhs_params, dst_params)) {
+    TFLITE_DCHECK(false);
+    return;
+  }
+
+  // Currently, only Ruy backend supports 16x8 quant gemm so we use ruy
+  // only.
+  detail::GemmImplUsingRuy<int8_t, int16_t, int32_t, int16_t,
+                           quantization_flavor>::Run(lhs_params, lhs_data,
+                                                     rhs_params, rhs_data,
+                                                     dst_params, dst_data,
+                                                     params, context);
+}
+
+// Special path for gemm with raw accumulator case. i.e. AccumScalar ==
+// DstScalar == int32 case.
+template <typename LhsScalar, typename RhsScalar,
+          QuantizationFlavor quantization_flavor>
+void Gemm(const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+          const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+          const MatrixParams<int32_t>& dst_params, int32_t* dst_data,
+          const GemmParams<int32_t, int32_t, quantization_flavor>& params,
+          CpuBackendContext* context) {
+  ruy::profiler::ScopeLabel label("cpu_backend_gemm::Gemm");
+  ValidateParams(lhs_params, rhs_params, dst_params, params);
+
+  // Currently, only Ruy backend supports get raw accumulator, so we use ruy
+  // only.
+  ruy::profiler::ScopeLabel label2("cpu_backend_gemm::Gemm: general GEMM");
+  detail::GemmImplUsingRuy<LhsScalar, RhsScalar, int32_t, int32_t,
+                           quantization_flavor>::Run(lhs_params, lhs_data,
+                                                     rhs_params, rhs_data,
+                                                     dst_params, dst_data,
+                                                     params, context);
+}
+
+}  // namespace cpu_backend_gemm
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_

tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h

@@ -0,0 +1,789 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Fast Gemv (i.e. matrix*vector multiplication) paths.
+// TODO(b/132094390): remove when GEMM performance is good enough on GEMV cases.
+
+// TFLite's runtime ops concentrate as much as possible the matrix*vector
+// use cases on the (matrix) * (column-vector) case, as opposed to
+// (row-vector) * (matrix).  So that is what we focus on optimizing here.
+// Accordingly, the public cpu_backend_gemm::Gemm() entry point checks
+// if we are in this (matrix) * (column-vector) case, and if so calls
+// CustomGemv.
+//
+// cpu_backend_gemm::Gemm is also currently restricted (as enforced in
+// ValidateParams) to the case where the left-hand side matrix is row-major.
+//
+// So the current scope of this CustomGemv function really is:
+// (row-major matrix) * (column-vector).
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_CUSTOM_GEMV_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_CUSTOM_GEMV_H_
+
+#include <stdint.h>
+
+#include <algorithm>
+#include <type_traits>
+#include <vector>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+// CustomGemvImpl is what needs to be specialized for each custom GEMV path.
+//
+// It does not deal with any multi-threaded implementation detail. Rather,
+// it provides the single-thread implementation to be run by each thread.
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct CustomGemvImpl {
+  // The number of rows of the left-hand-side matrix (and equivalently of the
+  // destination column-vector) that the kernel processes at a time.
+  // This will also be the minimum required number of rows for a Gemv shape
+  // to be supported by this path.
+  //
+  // Gemv implementations are expected to be able to deal with numbers of
+  // rows that aren't multiples of kKernelRows by possibly running the kernel
+  // again at an odd row_start, e.g. if kKernelRows==4, Run() should still
+  // support running on 7 rows by running twice: once with row_start=0 and then
+  // another time with row_start=3.
+  //
+  // On the other hand, gemv implementations are not expected to support
+  // running on fewer than kKernelRows rows. There is no interest in
+  // optimizing such narrow Gemv's that they are just a few dot-products.
+  // Supporting that would require custom kernel code only for that case.
+  static constexpr int kKernelRows = 1;
+
+  // Returns true if the Gemv shape is supported by Run(), provided that
+  // (row_end - row_start) > kKernelRows.
+  static bool IsSupportedGivenSufficientlyManyRows(
+      const MatrixParams<LhsScalar>& lhs_params,
+      const MatrixParams<RhsScalar>& rhs_params,
+      const MatrixParams<DstScalar>& dst_params,
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params) {
+    return false;
+  }
+
+  // Performs the Gemv.
+  static void Run(
+      const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+      const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+      int row_start, int row_end) {}
+};
+
+// Wraps CustomGemvImpl for multi-threaded operation.
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+class CustomGemvTask : public cpu_backend_threadpool::Task {
+ public:
+  CustomGemvTask(
+      const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+      const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+      int row_start, int row_end)
+      : lhs_params_(lhs_params),
+        lhs_data_(lhs_data),
+        rhs_params_(rhs_params),
+        rhs_data_(rhs_data),
+        dst_params_(dst_params),
+        dst_data_(dst_data),
+        params_(params),
+        row_start_(row_start),
+        row_end_(row_end) {}
+
+  void Run() override {
+    using Impl = CustomGemvImpl<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                                quantization_flavor>;
+    Impl::Run(lhs_params_, lhs_data_, rhs_params_, rhs_data_, dst_params_,
+              dst_data_, params_, row_start_, row_end_);
+  }
+
+ private:
+  const MatrixParams<LhsScalar>& lhs_params_;
+  const LhsScalar* lhs_data_;
+  const MatrixParams<RhsScalar>& rhs_params_;
+  const RhsScalar* rhs_data_;
+  const MatrixParams<DstScalar>& dst_params_;
+  DstScalar* dst_data_;
+  const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params_;
+  int row_start_;
+  int row_end_;
+};
+
+// Either performs the requested Gemv operation and returns true,
+// or immediately returns false.
+//
+// See the comment at the top of the file for the scope of what this handles.
+// In summary: (row-major matrix) * (column-vector).
+//
+// Here is only high-level logic.
+// The actual implementation details are in specializations of
+// CustomGemvImpl.
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+bool CustomGemv(
+    const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+    const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+    const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+    CpuBackendContext* context) {
+  ruy::profiler::ScopeLabel label("cpu_backend_gemm::Gemm: CustomGemv");
+  using Impl = CustomGemvImpl<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                              quantization_flavor>;
+  if (lhs_params.rows < Impl::kKernelRows) {
+    return false;
+  }
+  if (!Impl::IsSupportedGivenSufficientlyManyRows(lhs_params, rhs_params,
+                                                  dst_params, params)) {
+    return false;
+  }
+  TFLITE_DCHECK_GE(lhs_params.rows, Impl::kKernelRows);
+  int thread_count = LegacyHowManyThreads<Impl::kKernelRows>(
+      context->max_num_threads(), dst_params.rows, dst_params.cols,
+      lhs_params.cols);
+  if (thread_count == 1) {
+    Impl::Run(lhs_params, lhs_data, rhs_params, rhs_data, dst_params, dst_data,
+              params, 0, lhs_params.rows);
+  } else {
+    using Task = CustomGemvTask<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                                quantization_flavor>;
+    std::vector<Task> tasks;
+    tasks.reserve(thread_count);
+    const int kRowsPerThread =
+        RoundUp<Impl::kKernelRows>(CeilQuotient(dst_params.rows, thread_count));
+    int row_start = 0;
+    for (int i = 0; i < thread_count; i++) {
+      int row_end = std::min(dst_params.rows, row_start + kRowsPerThread);
+      tasks.emplace_back(lhs_params, lhs_data, rhs_params, rhs_data, dst_params,
+                         dst_data, params, row_start, row_end);
+      row_start = row_end;
+    }
+    cpu_backend_threadpool::Execute(tasks.size(), tasks.data(), context);
+  }
+  return true;
+}
+
+// USE_NEON still allows for x86 where we may be using the arm_neon_sse.h
+// wrapper implementing NEON intrinsics on top of SSE4 intrinsics.
+#ifdef USE_NEON
+
+// Some NEON helper functions used by CustomGemvImpl specializations below,
+// allowing for some type genericity in them.
+
+inline int16x8x2_t Load16AndSubtractZeroPoint(const std::uint8_t* src,
+                                              std::uint8_t zero_point) {
+  uint8x16_t src_u8 = vld1q_u8(src);
+  int16x8_t src_s16_0 = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(src_u8)));
+  int16x8_t src_s16_1 = vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(src_u8)));
+  int16x8x2_t result;
+  int16x8_t zero_point_vec = vdupq_n_s16(zero_point);
+  result.val[0] = vsubq_s16(src_s16_0, zero_point_vec);
+  result.val[1] = vsubq_s16(src_s16_1, zero_point_vec);
+  return result;
+}
+
+inline int16x8x2_t Load16AndSubtractZeroPoint(const std::int8_t* src,
+                                              std::int8_t zero_point) {
+  int8x16_t src_s8 = vld1q_s8(src);
+  int16x8_t src_s16_0 = vmovl_s8(vget_low_s8(src_s8));
+  int16x8_t src_s16_1 = vmovl_s8(vget_high_s8(src_s8));
+  int16x8x2_t result;
+  int16x8_t zero_point_vec = vdupq_n_s16(zero_point);
+  result.val[0] = vsubq_s16(src_s16_0, zero_point_vec);
+  result.val[1] = vsubq_s16(src_s16_1, zero_point_vec);
+  return result;
+}
+
+inline int16x8_t Load8AndSubtractZeroPoint(const std::uint8_t* src,
+                                           std::uint8_t zero_point) {
+  uint8x8_t src_u8 = vld1_u8(src);
+  int16x8_t src_s16 = vreinterpretq_s16_u16(vmovl_u8(src_u8));
+  int16x8_t zero_point_vec = vdupq_n_s16(zero_point);
+  return vsubq_s16(src_s16, zero_point_vec);
+}
+
+inline int16x8_t Load8AndSubtractZeroPoint(const std::int8_t* src,
+                                           std::int8_t zero_point) {
+  int8x8_t src_s8 = vld1_s8(src);
+  int16x8_t src_s16 = vmovl_s8(src_s8);
+  int16x8_t zero_point_vec = vdupq_n_s16(zero_point);
+  return vsubq_s16(src_s16, zero_point_vec);
+}
+
+inline void ClampAndStore(int32x4_t src, std::uint8_t clamp_min,
+                          std::uint8_t clamp_max, std::uint8_t* dst) {
+  // Narrow values down to 16 bit signed.
+  const int16x4_t res16 = vqmovn_s32(src);
+  // Narrow values down to 8 bit unsigned, saturating.
+  uint8x8_t res8 = vqmovun_s16(vcombine_s16(res16, res16));
+  // Apply the clamping from the activation function
+  res8 = vmax_u8(res8, vdup_n_u8(clamp_min));
+  res8 = vmin_u8(res8, vdup_n_u8(clamp_max));
+  // Store results to destination.
+  vst1_lane_u8(dst + 0, res8, 0);
+  vst1_lane_u8(dst + 1, res8, 1);
+  vst1_lane_u8(dst + 2, res8, 2);
+  vst1_lane_u8(dst + 3, res8, 3);
+}
+
+inline void ClampAndStore(int32x4_t src, std::int8_t clamp_min,
+                          std::int8_t clamp_max, std::int8_t* dst) {
+  // Narrow values down to 16 bit signed.
+  const int16x4_t res16 = vqmovn_s32(src);
+  // Narrow values down to 8 bit unsigned, saturating.
+  int8x8_t res8 = vqmovn_s16(vcombine_s16(res16, res16));
+  // Apply the clamping from the activation function
+  res8 = vmax_s8(res8, vdup_n_s8(clamp_min));
+  res8 = vmin_s8(res8, vdup_n_s8(clamp_max));
+  // Store results to destination.
+  vst1_lane_s8(dst + 0, res8, 0);
+  vst1_lane_s8(dst + 1, res8, 1);
+  vst1_lane_s8(dst + 2, res8, 2);
+  vst1_lane_s8(dst + 3, res8, 3);
+}
+
+inline void ClampAndStore(int32x4_t src, std::int16_t clamp_min,
+                          std::int16_t clamp_max, std::int16_t* dst) {
+  // Narrow values down to 16 bit signed.
+  int16x4_t res16 = vqmovn_s32(src);
+  // Apply the clamping from the activation function
+  res16 = vmax_s16(res16, vdup_n_s16(clamp_min));
+  res16 = vmin_s16(res16, vdup_n_s16(clamp_max));
+  // Store results to destination.
+  vst1_lane_s16(dst + 0, res16, 0);
+  vst1_lane_s16(dst + 1, res16, 1);
+  vst1_lane_s16(dst + 2, res16, 2);
+  vst1_lane_s16(dst + 3, res16, 3);
+}
+
+template <typename LhsScalar, typename RhsScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+struct CustomGemvImpl<LhsScalar, RhsScalar, std::int32_t, DstScalar,
+                      quantization_flavor> {
+  // This partial template specialization is less generic than its declaration
+  // implies: it assumes the following constraints on its free template
+  // parameters. We guard these assumptions in the following static_assert's.
+  static_assert(std::is_same<LhsScalar, std::uint8_t>::value ||
+                    std::is_same<LhsScalar, std::int8_t>::value,
+                "");
+  static_assert(std::is_same<RhsScalar, std::uint8_t>::value ||
+                    std::is_same<RhsScalar, std::int8_t>::value,
+                "");
+  static_assert(std::is_same<DstScalar, std::uint8_t>::value ||
+                    std::is_same<DstScalar, std::int8_t>::value ||
+                    std::is_same<DstScalar, std::int16_t>::value,
+                "");
+  static_assert(quantization_flavor ==
+                        QuantizationFlavor::kIntegerWithUniformMultiplier ||
+                    quantization_flavor ==
+                        QuantizationFlavor::kIntegerWithPerRowMultiplier,
+                "");
+
+  // This implementation's inner loop processes 4 rows of the left-hand side
+  // matrix at a time.
+  static constexpr int kKernelRows = 4;
+
+  static bool IsSupportedGivenSufficientlyManyRows(
+      const MatrixParams<LhsScalar>& lhs_params,
+      const MatrixParams<RhsScalar>& rhs_params,
+      const MatrixParams<DstScalar>& dst_params,
+      const GemmParams<std::int32_t, DstScalar, quantization_flavor>& params) {
+    // The kernel processes at least 8 LHS columns at once to fill NEON
+    // registers. The leftovers-handling code at the end works by loading a
+    // partially overlapping final register by walking back by a few (<8) values
+    // to avoid running past the row's end. This relies on there being
+    // at least 8 LHS columns.
+    return lhs_params.cols >= 8;
+  }
+
+  static void Run(
+      const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+      const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<std::int32_t, DstScalar, quantization_flavor>& params,
+      int row_start, int row_end) {
+    // Handle kKernelRows ( == 4) rows of the left-hand side matrix at each
+    // iteration of this for loop.
+    TFLITE_DCHECK_GE(row_end - row_start, kKernelRows);
+    for (int row = row_start; row < row_end; row += kKernelRows) {
+      // Here is the magic where we allow this kernel to handle any odd number
+      // of rows as long as it's >= kKernelRows: the last group of `kKernelRows`
+      // rows will be nudged to fit, possibly by starting at an odd value of
+      // `row`.
+      row = std::min(row, row_end - kKernelRows);
+      const LhsScalar* filter_ptr = lhs_data + row * lhs_params.cols;
+
+      static constexpr int kCacheLineSize = 64;
+      for (int k = 0; k < rhs_params.rows;
+           k += kCacheLineSize / sizeof(RhsScalar)) {
+        optimized_ops_preload_l1_keep(rhs_data + k);
+      }
+
+      // kPreloadAhead is empirically determined.
+      // End-to-end latency (ms) on mobilenet_v2_0.35_96_8bit, 1 thread,
+      // on Qualcomm S855:
+      //
+      // kPreloadAhead | big core | little core
+      // --------------+----------+------------
+      // 64            | 1.26     | 5.45
+      // 128           | 1.23     | 5.01
+      // 256           | 1.18     | 4.9
+      // 512           | 1.18     | 5.45
+      // 1024          | 1.18     | 6.5
+      // no prefetch   | 1.25     | 8.1
+      static constexpr int kPreloadAhead = 256;
+
+      // 4 accumulator registers, one for each row being processed.
+      // Each has 4 int32 lanes that corresponds to columns modulo 4, and
+      // will need to be horizontally reduced at the end.
+      int32x4_t acc0 = vdupq_n_s32(0);
+      int32x4_t acc1 = acc0;
+      int32x4_t acc2 = acc0;
+      int32x4_t acc3 = acc0;
+      int in = 0;
+      // As much as possible, handle 16 columns of the left-hand side matrix
+      // at a time. This allows for decent NEON implementation.
+      for (; in <= lhs_params.cols - 16; in += 16) {
+        const LhsScalar* local_filter_ptr = filter_ptr;
+        int16x8x2_t input_val =
+            Load16AndSubtractZeroPoint(rhs_data + in, rhs_params.zero_point);
+        int16x8x2_t filter_val_0 =
+            Load16AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(LhsScalar));
+        local_filter_ptr += lhs_params.cols;
+        int16x8x2_t filter_val_1 =
+            Load16AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(LhsScalar));
+        local_filter_ptr += lhs_params.cols;
+        int16x8x2_t filter_val_2 =
+            Load16AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(LhsScalar));
+        local_filter_ptr += lhs_params.cols;
+        int16x8x2_t filter_val_3 =
+            Load16AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(LhsScalar));
+        filter_ptr += 16;
+        acc0 = vmlal_s16(acc0, vget_low_s16(filter_val_0.val[0]),
+                         vget_low_s16(input_val.val[0]));
+        acc1 = vmlal_s16(acc1, vget_low_s16(filter_val_1.val[0]),
+                         vget_low_s16(input_val.val[0]));
+        acc2 = vmlal_s16(acc2, vget_low_s16(filter_val_2.val[0]),
+                         vget_low_s16(input_val.val[0]));
+        acc3 = vmlal_s16(acc3, vget_low_s16(filter_val_3.val[0]),
+                         vget_low_s16(input_val.val[0]));
+        acc0 = vmlal_s16(acc0, vget_low_s16(filter_val_0.val[1]),
+                         vget_low_s16(input_val.val[1]));
+        acc1 = vmlal_s16(acc1, vget_low_s16(filter_val_1.val[1]),
+                         vget_low_s16(input_val.val[1]));
+        acc2 = vmlal_s16(acc2, vget_low_s16(filter_val_2.val[1]),
+                         vget_low_s16(input_val.val[1]));
+        acc3 = vmlal_s16(acc3, vget_low_s16(filter_val_3.val[1]),
+                         vget_low_s16(input_val.val[1]));
+        acc0 = vmlal_s16(acc0, vget_high_s16(filter_val_0.val[0]),
+                         vget_high_s16(input_val.val[0]));
+        acc1 = vmlal_s16(acc1, vget_high_s16(filter_val_1.val[0]),
+                         vget_high_s16(input_val.val[0]));
+        acc2 = vmlal_s16(acc2, vget_high_s16(filter_val_2.val[0]),
+                         vget_high_s16(input_val.val[0]));
+        acc3 = vmlal_s16(acc3, vget_high_s16(filter_val_3.val[0]),
+                         vget_high_s16(input_val.val[0]));
+        acc0 = vmlal_s16(acc0, vget_high_s16(filter_val_0.val[1]),
+                         vget_high_s16(input_val.val[1]));
+        acc1 = vmlal_s16(acc1, vget_high_s16(filter_val_1.val[1]),
+                         vget_high_s16(input_val.val[1]));
+        acc2 = vmlal_s16(acc2, vget_high_s16(filter_val_2.val[1]),
+                         vget_high_s16(input_val.val[1]));
+        acc3 = vmlal_s16(acc3, vget_high_s16(filter_val_3.val[1]),
+                         vget_high_s16(input_val.val[1]));
+      }
+      // Less that 16 values remain. Try to handle 8 more.
+      if (in <= lhs_params.cols - 8) {
+        int16x8_t input_val =
+            Load8AndSubtractZeroPoint(rhs_data + in, rhs_params.zero_point);
+        int16x8_t filter_val_0 = Load8AndSubtractZeroPoint(
+            filter_ptr + 0 * lhs_params.cols, lhs_params.zero_point);
+        int16x8_t filter_val_1 = Load8AndSubtractZeroPoint(
+            filter_ptr + 1 * lhs_params.cols, lhs_params.zero_point);
+        int16x8_t filter_val_2 = Load8AndSubtractZeroPoint(
+            filter_ptr + 2 * lhs_params.cols, lhs_params.zero_point);
+        int16x8_t filter_val_3 = Load8AndSubtractZeroPoint(
+            filter_ptr + 3 * lhs_params.cols, lhs_params.zero_point);
+        filter_ptr += 8;
+        acc0 = vmlal_s16(acc0, vget_low_s16(filter_val_0),
+                         vget_low_s16(input_val));
+        acc1 = vmlal_s16(acc1, vget_low_s16(filter_val_1),
+                         vget_low_s16(input_val));
+        acc2 = vmlal_s16(acc2, vget_low_s16(filter_val_2),
+                         vget_low_s16(input_val));
+        acc3 = vmlal_s16(acc3, vget_low_s16(filter_val_3),
+                         vget_low_s16(input_val));
+        acc0 = vmlal_s16(acc0, vget_high_s16(filter_val_0),
+                         vget_high_s16(input_val));
+        acc1 = vmlal_s16(acc1, vget_high_s16(filter_val_1),
+                         vget_high_s16(input_val));
+        acc2 = vmlal_s16(acc2, vget_high_s16(filter_val_2),
+                         vget_high_s16(input_val));
+        acc3 = vmlal_s16(acc3, vget_high_s16(filter_val_3),
+                         vget_high_s16(input_val));
+        in += 8;
+      }
+      // Less than 8 values remain. Handle the remaining values
+      // in one more copy of the above code handling 8, where we
+      // walk back a few values to be able to load 8 values without
+      // overrunning the buffer. This is where we make use of the requirement
+      // (see IsSupportedGivenSufficientlyManyRows) that there at least
+      // 8 LHS columns.
+      if (in < lhs_params.cols) {
+        // `back` is how many entries to walk back by.
+        // Its value is necessarily between 1 and 7.
+        const int back = in + 8 - lhs_params.cols;
+        TFLITE_DCHECK_GE(back, 1);
+        TFLITE_DCHECK_LE(back, 7);
+        // Load 8 values as usual.
+        int16x8_t input_val = Load8AndSubtractZeroPoint(
+            rhs_data + lhs_params.cols - 8, rhs_params.zero_point);
+        const LhsScalar* local_filter_ptr = filter_ptr - back;
+        filter_ptr += lhs_params.cols - in;
+        int16x8_t filter_val_0 =
+            Load8AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        local_filter_ptr += lhs_params.cols;
+        int16x8_t filter_val_1 =
+            Load8AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        local_filter_ptr += lhs_params.cols;
+        int16x8_t filter_val_2 =
+            Load8AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        local_filter_ptr += lhs_params.cols;
+        int16x8_t filter_val_3 =
+            Load8AndSubtractZeroPoint(local_filter_ptr, lhs_params.zero_point);
+        // Now zero out the `back` first entries of input_val.
+        // vsetq_lane_s16 takes a literal index, so we need unrolled code.
+        switch (back) {
+          case 7:
+            input_val = vsetq_lane_s16(0, input_val, 6);
+            [[clang::fallthrough]];
+          case 6:
+            input_val = vsetq_lane_s16(0, input_val, 5);
+            [[clang::fallthrough]];
+          case 5:
+            input_val = vsetq_lane_s16(0, input_val, 4);
+            [[clang::fallthrough]];
+          case 4:
+            input_val = vsetq_lane_s16(0, input_val, 3);
+            [[clang::fallthrough]];
+          case 3:
+            input_val = vsetq_lane_s16(0, input_val, 2);
+            [[clang::fallthrough]];
+          case 2:
+            input_val = vsetq_lane_s16(0, input_val, 1);
+            [[clang::fallthrough]];
+          default:
+            input_val = vsetq_lane_s16(0, input_val, 0);
+        }
+        // Multiply-accumulate 8 values as usual. The `back` first lanes
+        // of filter_val_* are junk, but it doesn't matter since they get
+        // multiplied by the zeros that we just wrote in the corresponding
+        // lanes of input_val.
+        acc0 = vmlal_s16(acc0, vget_low_s16(filter_val_0),
+                         vget_low_s16(input_val));
+        acc1 = vmlal_s16(acc1, vget_low_s16(filter_val_1),
+                         vget_low_s16(input_val));
+        acc2 = vmlal_s16(acc2, vget_low_s16(filter_val_2),
+                         vget_low_s16(input_val));
+        acc3 = vmlal_s16(acc3, vget_low_s16(filter_val_3),
+                         vget_low_s16(input_val));
+        acc0 = vmlal_s16(acc0, vget_high_s16(filter_val_0),
+                         vget_high_s16(input_val));
+        acc1 = vmlal_s16(acc1, vget_high_s16(filter_val_1),
+                         vget_high_s16(input_val));
+        acc2 = vmlal_s16(acc2, vget_high_s16(filter_val_2),
+                         vget_high_s16(input_val));
+        acc3 = vmlal_s16(acc3, vget_high_s16(filter_val_3),
+                         vget_high_s16(input_val));
+      }
+
+      // Horizontally reduce accumulators
+      int32x2_t pairwise_reduced_acc_0 =
+          vpadd_s32(vget_low_s32(acc0), vget_high_s32(acc0));
+      int32x2_t pairwise_reduced_acc_1 =
+          vpadd_s32(vget_low_s32(acc1), vget_high_s32(acc1));
+      int32x2_t pairwise_reduced_acc_2 =
+          vpadd_s32(vget_low_s32(acc2), vget_high_s32(acc2));
+      int32x2_t pairwise_reduced_acc_3 =
+          vpadd_s32(vget_low_s32(acc3), vget_high_s32(acc3));
+      const int32x2_t reduced_lo =
+          vpadd_s32(pairwise_reduced_acc_0, pairwise_reduced_acc_1);
+      const int32x2_t reduced_hi =
+          vpadd_s32(pairwise_reduced_acc_2, pairwise_reduced_acc_3);
+      int32x4_t reduced = vcombine_s32(reduced_lo, reduced_hi);
+      // End of horizontal reduction: now `reduced` is a single int32x4
+      // containing the 4 int32 accumulators corresponding to the 4 rows
+      // being processed.
+
+      // Add bias values.
+      if (params.bias) {
+        int32x4_t bias_vec = vld1q_s32(params.bias + row);
+        reduced = vaddq_s32(reduced, bias_vec);
+      }
+
+      // Get multiplier parameters.
+      int32x4_t multiplier_fixedpoint;
+      int32x4_t multiplier_exponent;
+      if (quantization_flavor ==
+          QuantizationFlavor::kIntegerWithPerRowMultiplier) {
+        multiplier_exponent =
+            vld1q_s32(params.multiplier_exponent_perchannel + row);
+        multiplier_fixedpoint =
+            vld1q_s32(params.multiplier_fixedpoint_perchannel + row);
+      } else {
+        multiplier_exponent = vdupq_n_s32(params.multiplier_exponent);
+        multiplier_fixedpoint = vdupq_n_s32(params.multiplier_fixedpoint);
+      }
+
+      // If positive exponent, shift left.
+      int32x4_t exponent_positive_part =
+          vmaxq_s32(multiplier_exponent, vdupq_n_s32(0));
+      reduced = vshlq_s32(reduced, exponent_positive_part);
+      // Multiply by the fixed-point multiplier.
+      reduced = vqrdmulhq_s32(reduced, multiplier_fixedpoint);
+      // If negative exponent, rounding-shift-right.
+      int32x4_t exponent_negative_part =
+          vminq_s32(multiplier_exponent, vdupq_n_s32(0));
+      reduced = vrshlq_s32(reduced, exponent_negative_part);
+
+      // Add the output offset.
+      const int32x4_t output_offset_vec = vdupq_n_s32(dst_params.zero_point);
+      reduced = vaddq_s32(reduced, output_offset_vec);
+
+      // Finally, clamp and store to the destination.
+      ClampAndStore(reduced, params.clamp_min, params.clamp_max,
+                    dst_data + row);
+    }
+  }
+};
+
+// The float specialization below is unconditionally faster than ruy
+// because ruy does not currently have any Gemv path.
+// But it is not unconditionally faster than Eigen, which is what is used
+// unless TFLITE_WITH_RUY is defined. Indeed, Eigen has decently efficient
+// Gemv paths, and they may use AVX instructions, while the present
+// NEON intrinsics code maps at best to SSE4 on x86.
+#ifdef TFLITE_WITH_RUY
+
+// We want to use fused multiply-add when it's available (that is, on A64
+// unconditionally and on A32 with VFPv4) because it's often faster, and
+// because non-fused seems not to be available in A64 so a conscientious
+// compiler might emit slow code (separate mul and add instructions) in order to
+// implement the vmlaq_f32 intrinsic with strict bit-for-bit exactness on A64.
+// (Compilers seem to be generating a fused fmla instruction at the moment,
+// but that could change).
+//
+// We still want to support building for A32 without VFPv4.
+inline float32x4_t mul_add(float32x4_t acc, float32x4_t lhs, float32x4_t rhs) {
+#ifdef __ARM_FEATURE_FMA
+  return vfmaq_f32(acc, lhs, rhs);
+#else
+  return vmlaq_f32(acc, lhs, rhs);
+#endif
+}
+
+template <>
+struct CustomGemvImpl<float, float, float, float,
+                      QuantizationFlavor::kFloatingPoint> {
+  // This implementation's inner loop processes 4 rows of the left-hand side
+  // matrix at a time.
+  static constexpr int kKernelRows = 4;
+
+  static bool IsSupportedGivenSufficientlyManyRows(
+      const MatrixParams<float>& lhs_params,
+      const MatrixParams<float>& rhs_params,
+      const MatrixParams<float>& dst_params,
+      const GemmParams<float, float>& params) {
+    // The kernel processes 4 LHS columns at once to fill float32x4 registers.
+    // The leftovers-handling code at the end works by loading a partially
+    // overlapping final register by walking back by a few (<4) floats
+    // to avoid running past the row's end. This relies on there being
+    // at least 4 LHS columns.
+    return lhs_params.cols >= 4;
+  }
+  static void Run(const MatrixParams<float>& lhs_params, const float* lhs_data,
+                  const MatrixParams<float>& rhs_params, const float* rhs_data,
+                  const MatrixParams<float>& dst_params, float* dst_data,
+                  const GemmParams<float, float>& params, int row_start,
+                  int row_end) {
+    // Handle kKernelRows ( == 4) rows of the left-hand side matrix at each
+    // iteration of this for loop.
+    TFLITE_DCHECK_GE(row_end - row_start, kKernelRows);
+    for (int row = row_start; row < row_end; row += kKernelRows) {
+      // Here is the magic where we allow this kernel to handle any odd number
+      // of rows as long as it's >= kKernelRows: the last group of `kKernelRows`
+      // rows will be nudged to fit, possibly by starting at an odd value of
+      // `row`.
+      row = std::min(row, row_end - kKernelRows);
+      const float* filter_ptr = lhs_data + row * lhs_params.cols;
+
+      static constexpr int kCacheLineSize = 64;
+      for (int k = 0; k < rhs_params.rows;
+           k += kCacheLineSize / sizeof(float)) {
+        optimized_ops_preload_l1_keep(rhs_data + k);
+      }
+
+      // kPreloadAhead is empirically determined.
+      // End-to-end latency (ms) on mobilenet_v2_0.35_96_float, 1 thread,
+      // on Qualcomm S855:
+      //
+      // kPreloadAhead | big core | little core
+      // --------------+----------+------------
+      // 64            | 2.4      | 15.2
+      // 128           | 2.15     | 12.9
+      // 256           | 2        | 12.9
+      // 512           | 2.08     | 13.3
+      // 1024          | 2.05     | 14.7
+      // no prefetch   | 2.1      | 28
+      static constexpr int kPreloadAhead = 256;
+
+      // 4 accumulator registers, one for each row being processed.
+      // Each has 4 float32 lanes that corresponds to columns modulo 4, and
+      // will need to be horizontally reduced at the end.
+      float32x4_t acc0 = vdupq_n_f32(0);
+      float32x4_t acc1 = acc0;
+      float32x4_t acc2 = acc0;
+      float32x4_t acc3 = acc0;
+      int in = 0;
+      // As much as possible, handle 4 columns of the left-hand side matrix
+      // at a time. This allows for decent NEON implementation.
+      for (; in <= lhs_params.cols - 4; in += 4) {
+        float32x4_t input_val = vld1q_f32(rhs_data + in);
+        const float* local_filter_ptr = filter_ptr;
+        float32x4_t filter_val_0 = vld1q_f32(local_filter_ptr);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(float));
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_1 = vld1q_f32(local_filter_ptr);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(float));
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_2 = vld1q_f32(local_filter_ptr);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(float));
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_3 = vld1q_f32(local_filter_ptr);
+        optimized_ops_preload_l1_stream(local_filter_ptr +
+                                        kPreloadAhead / sizeof(float));
+        filter_ptr += 4;
+        acc0 = mul_add(acc0, filter_val_0, input_val);
+        acc1 = mul_add(acc1, filter_val_1, input_val);
+        acc2 = mul_add(acc2, filter_val_2, input_val);
+        acc3 = mul_add(acc3, filter_val_3, input_val);
+      }
+      // Less than 4 values remain. Handle the remaining values
+      // in one more copy of the above code handling 4, where we
+      // walk back a few values to be able to load 4 values without
+      // overrunning the buffer. This is where we make use of the requirement
+      // (see IsSupportedGivenSufficientlyManyRows) that there at least
+      // 4 LHS columns.
+      if (in < lhs_params.cols) {
+        // `back` is how many entries to walk back by.
+        // Its value is necessarily between 1 and 3.
+        const int back = in + 4 - lhs_params.cols;
+        TFLITE_DCHECK_GE(back, 1);
+        TFLITE_DCHECK_LE(back, 3);
+        // Load 4 values as usual.
+        float32x4_t input_val = vld1q_f32(rhs_data + lhs_params.cols - 4);
+        const float* local_filter_ptr = filter_ptr - back;
+        filter_ptr += lhs_params.cols - in;
+        float32x4_t filter_val_0 = vld1q_f32(local_filter_ptr);
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_1 = vld1q_f32(local_filter_ptr);
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_2 = vld1q_f32(local_filter_ptr);
+        local_filter_ptr += lhs_params.cols;
+        float32x4_t filter_val_3 = vld1q_f32(local_filter_ptr);
+        // Now zero out the `back` first entries of input_val.
+        // vsetq_lane_f32 takes a literal index, so we need unrolled code.
+        switch (back) {
+          case 3:
+            input_val = vsetq_lane_f32(0, input_val, 2);
+            [[clang::fallthrough]];
+          case 2:
+            input_val = vsetq_lane_f32(0, input_val, 1);
+            [[clang::fallthrough]];
+          default:
+            input_val = vsetq_lane_f32(0, input_val, 0);
+        }
+        // Multiply-accumulate 4 values as usual. The `back` first lanes
+        // of filter_val_* are junk, but it doesn't matter since they get
+        // multiplied by the zeros that we just wrote in the corresponding
+        // lanes of input_val.
+        acc0 = mul_add(acc0, filter_val_0, input_val);
+        acc1 = mul_add(acc1, filter_val_1, input_val);
+        acc2 = mul_add(acc2, filter_val_2, input_val);
+        acc3 = mul_add(acc3, filter_val_3, input_val);
+      }
+
+      // Horizontally reduce accumulators
+      float32x2_t pairwise_reduced_acc_0 =
+          vpadd_f32(vget_low_f32(acc0), vget_high_f32(acc0));
+      float32x2_t pairwise_reduced_acc_1 =
+          vpadd_f32(vget_low_f32(acc1), vget_high_f32(acc1));
+      float32x2_t pairwise_reduced_acc_2 =
+          vpadd_f32(vget_low_f32(acc2), vget_high_f32(acc2));
+      float32x2_t pairwise_reduced_acc_3 =
+          vpadd_f32(vget_low_f32(acc3), vget_high_f32(acc3));
+      float32x2_t reduced_lo =
+          vpadd_f32(pairwise_reduced_acc_0, pairwise_reduced_acc_1);
+      float32x2_t reduced_hi =
+          vpadd_f32(pairwise_reduced_acc_2, pairwise_reduced_acc_3);
+      float32x4_t reduced = vcombine_f32(reduced_lo, reduced_hi);
+      // End of horizontal reduction: now `reduced` is a single float32x4
+      // containing the 4 float32 accumulators corresponding to the 4 rows
+      // being processed.
+
+      if (params.bias) {
+        // Add bias values.
+        reduced = vaddq_f32(reduced, vld1q_f32(params.bias + row));
+      }
+
+      // Clamp and store to destination.
+      reduced = vminq_f32(reduced, vdupq_n_f32(params.clamp_max));
+      reduced = vmaxq_f32(reduced, vdupq_n_f32(params.clamp_min));
+      vst1q_f32(dst_data + row, reduced);
+    }
+  }
+};
+
+#endif  // TFLITE_WITH_RUY
+
+#endif  // USE_NEON
+
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_CUSTOM_GEMV_H_

tensorflow/lite/kernels/cpu_backend_gemm_eigen.cc

@@ -0,0 +1,81 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TFLITE_WITH_RUY
+
+#include "tensorflow/lite/kernels/cpu_backend_gemm_eigen.h"
+
+// See b/131835803: in TFLite code, because eigen_spatial_convolutions.h does
+// #define Eigen EigenForTFLite, it is difficult to have any #include of Eigen
+// headers in a header file, as that results in name classes (compilation
+// errors) depending on the order in which these headers are #included.
+// So we have moved the #include of Eigen here, in a .cc file, where we have
+// control over the header #include sequence.
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+// This function is out-of-line in a .cc file because of the issue
+// noted above in the comment on the #include for Eigen/Core.
+void GemmImplUsingEigen::Run(
+    const MatrixParams<float>& lhs_params, const float* lhs_data,
+    const MatrixParams<float>& rhs_params, const float* rhs_data,
+    const MatrixParams<float>& dst_params, float* dst_data,
+    const GemmParams<float, float>& params, CpuBackendContext* /* context */) {
+  // This code assumes specific storage orders, encoded in these Eigen types.
+  // These assumptions have been checked by TF_LITE_ASSERT's in the public
+  // Gemm entry point already, before the implementation gets to this point.
+  using EigenMatrixMapRowMajorConst =
+      Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic,
+                                     Eigen::RowMajor>>;
+  using EigenMatrixMapColMajorConst =
+      Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic,
+                                     Eigen::ColMajor>>;
+  using EigenMatrixMapColMajorMutable = Eigen::Map<
+      Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>>;
+
+  EigenMatrixMapRowMajorConst eigen_lhs(lhs_data, lhs_params.rows,
+                                        lhs_params.cols);
+  EigenMatrixMapColMajorConst eigen_rhs(rhs_data, rhs_params.rows,
+                                        rhs_params.cols);
+  EigenMatrixMapColMajorMutable eigen_dst(dst_data, dst_params.rows,
+                                          dst_params.cols);
+
+  if (rhs_params.cols == 1) {
+    eigen_dst.col(0).noalias() = eigen_lhs * eigen_rhs.col(0);
+  } else if (lhs_params.rows == 1) {
+    eigen_dst.row(0).noalias() = eigen_lhs.row(0) * eigen_rhs;
+  } else {
+    eigen_dst.noalias() = eigen_lhs * eigen_rhs;
+  }
+
+  if (params.bias) {
+    BiasAndClamp(params.clamp_min, params.clamp_max, dst_params.rows,
+                 params.bias, dst_params.rows * dst_params.cols, dst_data);
+  } else {
+    eigen_dst = eigen_dst.cwiseMin(params.clamp_max).cwiseMax(params.clamp_min);
+  }
+}
+
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // not TFLITE_WITH_RUY

tensorflow/lite/kernels/cpu_backend_gemm_eigen.h

@@ -0,0 +1,42 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_EIGEN_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_EIGEN_H_
+
+#ifndef TFLITE_WITH_RUY
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+struct GemmImplUsingEigen {
+  static void Run(const MatrixParams<float>& lhs_params, const float* lhs_data,
+                  const MatrixParams<float>& rhs_params, const float* rhs_data,
+                  const MatrixParams<float>& dst_params, float* dst_data,
+                  const GemmParams<float, float>& params,
+                  CpuBackendContext* /* context */);
+};
+
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // not TFLITE_WITH_RUY
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_EIGEN_H_

tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h

@@ -0,0 +1,195 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_GEMMLOWP_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_GEMMLOWP_H_
+
+#include <tuple>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#ifndef TFLITE_WITH_RUY
+
+#include <cstdint>
+#include <type_traits>
+
+#include "public/gemmlowp.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+template <typename DstScalar>
+struct GemmlowpSaturatingCastStage {};
+
+template <>
+struct GemmlowpSaturatingCastStage<std::uint8_t> {
+  using Type = gemmlowp::OutputStageSaturatingCastToUint8;
+};
+
+template <>
+struct GemmlowpSaturatingCastStage<std::int8_t> {
+  using Type = gemmlowp::OutputStageSaturatingCastToInt8;
+};
+
+template <>
+struct GemmlowpSaturatingCastStage<std::int16_t> {
+  using Type = gemmlowp::OutputStageSaturatingCastToInt16;
+};
+
+template <typename DstScalar>
+struct GemmlowpBitDepthParams {};
+
+template <>
+struct GemmlowpBitDepthParams<std::uint8_t> {
+  using Type = gemmlowp::L8R8WithLhsNonzeroBitDepthParams;
+};
+
+template <>
+struct GemmlowpBitDepthParams<std::int8_t> {
+  using Type = gemmlowp::SignedL8R8WithLhsNonzeroBitDepthParams;
+};
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImplUsingGemmlowp {};
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+struct GemmImplUsingGemmlowp<
+    LhsScalar, RhsScalar, AccumScalar, DstScalar,
+    QuantizationFlavor::kIntegerWithUniformMultiplier> {
+  static_assert(std::is_same<LhsScalar, RhsScalar>::value, "");
+  static_assert(std::is_same<AccumScalar, std::int32_t>::value, "");
+  using SrcScalar = LhsScalar;
+
+  static void Run(
+      const MatrixParams<SrcScalar>& lhs_params, const SrcScalar* lhs_data,
+      const MatrixParams<SrcScalar>& rhs_params, const SrcScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<std::int32_t, DstScalar,
+                       QuantizationFlavor::kIntegerWithUniformMultiplier>&
+          params,
+      CpuBackendContext* context) {
+    gemmlowp::MatrixMap<const SrcScalar, gemmlowp::MapOrder::RowMajor>
+        gemmlowp_lhs(lhs_data, lhs_params.rows, lhs_params.cols);
+    gemmlowp::MatrixMap<const SrcScalar, gemmlowp::MapOrder::ColMajor>
+        gemmlowp_rhs(rhs_data, rhs_params.rows, rhs_params.cols);
+    gemmlowp::MatrixMap<DstScalar, gemmlowp::MapOrder::ColMajor> gemmlowp_dst(
+        dst_data, dst_params.rows, dst_params.cols);
+
+    using ColVectorMap =
+        gemmlowp::VectorMap<const int32_t, gemmlowp::VectorShape::Col>;
+    gemmlowp::OutputStageScaleInt32ByFixedPointAndExponent scale_stage;
+    scale_stage.result_offset_after_shift = dst_params.zero_point;
+    scale_stage.result_fixedpoint_multiplier = params.multiplier_fixedpoint;
+    scale_stage.result_exponent = params.multiplier_exponent;
+    using SaturatingCastStageType =
+        typename GemmlowpSaturatingCastStage<DstScalar>::Type;
+    gemmlowp::OutputStageClamp clamp_stage;
+    clamp_stage.min = params.clamp_min;
+    clamp_stage.max = params.clamp_max;
+    SaturatingCastStageType saturating_cast_stage;
+    using BitDepthParams = typename GemmlowpBitDepthParams<SrcScalar>::Type;
+    if (params.bias) {
+      ColVectorMap bias_vector(params.bias, lhs_params.rows);
+      gemmlowp::OutputStageBiasAddition<ColVectorMap> bias_addition_stage;
+      bias_addition_stage.bias_vector = bias_vector;
+      auto output_pipeline = std::make_tuple(
+          bias_addition_stage, scale_stage, clamp_stage, saturating_cast_stage);
+      gemmlowp::GemmWithOutputPipeline<SrcScalar, DstScalar, BitDepthParams>(
+          context->gemmlowp_context(), gemmlowp_lhs, gemmlowp_rhs,
+          &gemmlowp_dst, -lhs_params.zero_point, -rhs_params.zero_point,
+          output_pipeline);
+    } else {
+      auto output_pipeline =
+          std::make_tuple(scale_stage, clamp_stage, saturating_cast_stage);
+      gemmlowp::GemmWithOutputPipeline<SrcScalar, DstScalar, BitDepthParams>(
+          context->gemmlowp_context(), gemmlowp_lhs, gemmlowp_rhs,
+          &gemmlowp_dst, -lhs_params.zero_point, -rhs_params.zero_point,
+          output_pipeline);
+    }
+  }
+};
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+struct GemmImplUsingGemmlowp<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                             QuantizationFlavor::kIntegerWithPerRowMultiplier> {
+  static_assert(std::is_same<LhsScalar, RhsScalar>::value, "");
+  static_assert(std::is_same<AccumScalar, std::int32_t>::value, "");
+  using SrcScalar = LhsScalar;
+
+  static void Run(
+      const MatrixParams<SrcScalar>& lhs_params, const SrcScalar* lhs_data,
+      const MatrixParams<SrcScalar>& rhs_params, const SrcScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<std::int32_t, DstScalar,
+                       QuantizationFlavor::kIntegerWithPerRowMultiplier>&
+          params,
+      CpuBackendContext* context) {
+    // gemmlowp support for this per-channel path is limited to NEON.
+    // We fall back to ruy outside of NEON.
+#ifdef GEMMLOWP_NEON
+    gemmlowp::MatrixMap<const SrcScalar, gemmlowp::MapOrder::RowMajor>
+        gemmlowp_lhs(lhs_data, lhs_params.rows, lhs_params.cols);
+    gemmlowp::MatrixMap<const SrcScalar, gemmlowp::MapOrder::ColMajor>
+        gemmlowp_rhs(rhs_data, rhs_params.rows, rhs_params.cols);
+    gemmlowp::MatrixMap<DstScalar, gemmlowp::MapOrder::ColMajor> gemmlowp_dst(
+        dst_data, dst_params.rows, dst_params.cols);
+
+    using ColVectorMap =
+        gemmlowp::VectorMap<const int32, gemmlowp::VectorShape::Col>;
+    ColVectorMap bias_vector(params.bias, lhs_params.rows);
+    gemmlowp::OutputStageBiasAddition<ColVectorMap> bias_addition_stage;
+    bias_addition_stage.bias_vector = bias_vector;
+    gemmlowp::OutputStageScaleInt32ByFixedPointAndExponentPC<
+        gemmlowp::VectorShape::Col>
+        scale_stage;
+    scale_stage.result_offset_after_shift = dst_params.zero_point;
+    scale_stage.result_fixedpoint_multiplier =
+        ColVectorMap(params.multiplier_fixedpoint_perchannel, dst_params.rows);
+    scale_stage.result_exponent =
+        ColVectorMap(params.multiplier_exponent_perchannel, dst_params.rows);
+    using SaturatingCastStageType =
+        typename GemmlowpSaturatingCastStage<DstScalar>::Type;
+    gemmlowp::OutputStageClamp clamp_stage;
+    clamp_stage.min = params.clamp_min;
+    clamp_stage.max = params.clamp_max;
+    SaturatingCastStageType saturating_cast_stage;
+    auto output_pipeline = std::make_tuple(bias_addition_stage, scale_stage,
+                                           clamp_stage, saturating_cast_stage);
+    using BitDepthParams = typename GemmlowpBitDepthParams<SrcScalar>::Type;
+    gemmlowp::GemmWithOutputPipeline<SrcScalar, DstScalar, BitDepthParams>(
+        context->gemmlowp_context(), gemmlowp_lhs, gemmlowp_rhs, &gemmlowp_dst,
+        -lhs_params.zero_point, -rhs_params.zero_point, output_pipeline);
+#else
+    GemmImplUsingRuy<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                     QuantizationFlavor::kIntegerWithPerRowMultiplier>::
+        Run(lhs_params, lhs_data, rhs_params, rhs_data, dst_params, dst_data,
+            params, context);
+#endif
+  }
+};
+
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // not TFLITE_WITH_RUY
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_GEMMLOWP_H_

tensorflow/lite/kernels/cpu_backend_gemm_params.h

@@ -0,0 +1,263 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_PARAMS_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_PARAMS_H_
+
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+
+namespace cpu_backend_gemm {
+
+// Matrix storage order: column-major or row-major.
+enum class Order { kColMajor, kRowMajor };
+
+enum class CachePolicy : std::uint8_t {
+  kNeverCache,
+  kCacheIfLargeSpeedup,
+  kAlwaysCache,
+};
+
+inline CachePolicy DefaultCachePolicy(bool is_constant_data) {
+  return is_constant_data ? CachePolicy::kCacheIfLargeSpeedup
+                          : CachePolicy::kNeverCache;
+}
+
+// MatrixParams encapsulates the parameters that Gemm needs about each
+// matrix, besides the buffer data pointer.
+// Compare to ruy::Matrix, which also encapsulates the data pointer.
+// Rationale for leaving the data pointer out of here: doing so
+// requires complicated const-correctness mechanics. See
+// ruy::ConstCheckingPtr.
+template <typename Scalar>
+struct MatrixParams {
+  // Storage layout order. For now we only do plain linear non-strided
+  // layout. It would be easy to support a stride if needed.
+  Order order = Order::kColMajor;
+  // Number of rows of the matrix.
+  int rows = 0;
+  // Number of columns of the matrix.
+  int cols = 0;
+  // The zero_point, i.e. which Scalar value is to be interpreted as zero.
+  // When Scalar is floating-point, this must be 0.
+  Scalar zero_point = 0;
+  // When the data pointed to by this matrix is constant data, so that it is
+  // valid to assume that equality of pointers implies equality of data,
+  // a CachePolicy may be used instead of the default kNeverCache,
+  // which will enable ruy to take advantage of this constancy of the data to
+  // cache the packing work, which can be a large speedup in matrix*vector
+  // and other narrow shapes.
+  CachePolicy cache_policy = CachePolicy::kNeverCache;
+};
+
+// Enumeration of broad categories of Gemm.
+//
+// The primary reason for this to exist is to allow Gemm to compile
+// only uniform-quantized or only per-channel-quantized code paths.
+// This is unneeded with ruy as the back-end, as this is only a runtime
+// difference in ruy, but with gemmlowp these really are separate code
+// paths and templatizing in a QuantizationFlavor is necessary to avoid
+// compiling unused gemmlowp code. Indeed, TFLite currently uses
+// uint8 with uniform quantization and int8 with per-channel quantization,
+// and does not use uint8 with per-channel. We want to avoid compiling
+// the gemmlowp uint8 per-channel path when gemmlowp is the back-end.
+//
+// It's possible to drop this in the future if gemmlowp goes away and no
+// other then-relevant backend library handles quantized paths in a way that
+// requires knowing this at compile-time.
+enum class QuantizationFlavor {
+  // Floating-point Gemm: the accumulators are not multiplied by any
+  // 'multiplier'.
+  kFloatingPoint,
+  // Quantized Gemm using a single multiplier for all accumulators.
+  kIntegerWithUniformMultiplier,
+  // Quantized Gemm using a separate multipliers for accumulators of each
+  // row of the destination matrix. This is what is called 'per-channel'
+  // in GemmParams. Here we use the more specific 'per-row' terminology
+  // to allow for the possibility of 'per-column' in the future, and to
+  // allow for that to be a separate code path in some back-end such as
+  // gemmlowp.
+  kIntegerWithPerRowMultiplier
+};
+
+// Additional parameters that Gemm needs, beyond what falls into
+// the MatrixParams that it takes. Compare to ruy::Spec.
+//
+// Decoupling AccumScalar from DstScalar (rather than deducing it from that)
+// is useful future-proofing. Think of a float16 path using float32 accum.
+//
+// QuantizationFlavor is passed here even though it's technically not used
+// in this class. This is so that we retain the ability in the future to
+// specialize this class for quantization flavor, and this allows for
+// Gemm to be templatized in quantization_flavor via the GemmParams that it
+// takes, allowing for automatic template parameter deduction to take place,
+// so that most call sites don't need to specify a QuantizationFlavor
+// (only those that need perchannel quantization do).
+template <typename AccumScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor =
+              std::is_floating_point<AccumScalar>::value
+                  ? QuantizationFlavor::kFloatingPoint
+                  : QuantizationFlavor::kIntegerWithUniformMultiplier>
+struct GemmParams {
+  // Only for non-floating-point cases. The fixed-point part (i.e. the mantissa)
+  // of the multiplier by which accumulators are multiplied before being casted
+  // to the destination type.
+  AccumScalar multiplier_fixedpoint = 0;
+  // Only for non-floating-point cases. The exponent part of the aforementioned
+  // multiplier.
+  int multiplier_exponent = 0;
+  // Per-channel variant of multiplier_fixedpoint. If not nullptr, this must
+  // point to a buffer of as many values as there are rows in the destination
+  // matrix. Each row of the destination matrix will use the corresponding
+  // buffer element instead of multiplier_fixedpoint.
+  const AccumScalar* multiplier_fixedpoint_perchannel = nullptr;
+  // Per-channel variant of multiplier_exponent. If not nullptr, this must
+  // point to a buffer of as many values as there are rows in the destination
+  // matrix. Each row of the destination matrix will use the corresponding
+  // buffer element instead of multiplier_exponent.
+  //
+  // Either none or both of multiplier_exponent_perchannel and
+  // multiplier_fixedpoint_perchannel must be nullptr.
+  const int* multiplier_exponent_perchannel = nullptr;
+  // The bias vector data, if not null.
+  const AccumScalar* bias = nullptr;
+  // min clamp bound of destination values.
+  DstScalar clamp_min = std::is_floating_point<DstScalar>::value
+                            ? -std::numeric_limits<DstScalar>::infinity()
+                            : std::numeric_limits<DstScalar>::lowest();
+  // max clamp bound of destination values.
+  DstScalar clamp_max = std::is_floating_point<DstScalar>::value
+                            ? std::numeric_limits<DstScalar>::infinity()
+                            : std::numeric_limits<DstScalar>::max();
+};
+
+/* Convenience typedefs */
+
+template <typename DstScalar>
+using QuantizedGemmParams = GemmParams<std::int32_t, DstScalar>;
+
+using FloatGemmParams = GemmParams<float, float>;
+
+/* Validation functions */
+
+// Note that this uses TFLITE_DCHECK from kernels/internal/compatibility.h
+// and not TF_LITE_ASSERT from op_macros.h. We want this to be explicitly
+// debug-build-only assertions so that there's not reason not to
+// generously validate, and TF_LITE_ASSERT is actually at the moment
+// a release-build assertion. See b/131587258.
+
+// Validates self-consistency of GemmParams.
+template <typename AccumScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+void ValidateGemmParams(
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params) {
+  // Guard consistency of the quantized multiplier fields.
+  if (quantization_flavor == QuantizationFlavor::kFloatingPoint) {
+    TFLITE_DCHECK(!params.multiplier_fixedpoint);
+    TFLITE_DCHECK(!params.multiplier_exponent);
+    TFLITE_DCHECK(!params.multiplier_fixedpoint_perchannel);
+    TFLITE_DCHECK(!params.multiplier_exponent_perchannel);
+  } else if (quantization_flavor ==
+                 QuantizationFlavor::kIntegerWithUniformMultiplier &&
+             !std::is_same<DstScalar, int32_t>::value) {
+    TFLITE_DCHECK(params.multiplier_fixedpoint);
+    // Nothing to check about multiplier_exponent
+    TFLITE_DCHECK(!params.multiplier_fixedpoint_perchannel);
+    TFLITE_DCHECK(!params.multiplier_exponent_perchannel);
+  } else if (quantization_flavor ==
+                 QuantizationFlavor::kIntegerWithPerRowMultiplier &&
+             !std::is_same<DstScalar, int32_t>::value) {
+    TFLITE_DCHECK(!params.multiplier_fixedpoint);
+    TFLITE_DCHECK(!params.multiplier_exponent);
+    TFLITE_DCHECK(params.multiplier_fixedpoint_perchannel);
+    TFLITE_DCHECK(params.multiplier_exponent_perchannel);
+  } else {
+    // For the get raw accumulator case, we should make sure none of the
+    // quantization params are set.
+    TFLITE_DCHECK(!params.multiplier_fixedpoint);
+    TFLITE_DCHECK(!params.multiplier_exponent);
+    TFLITE_DCHECK(!params.multiplier_fixedpoint_perchannel);
+    TFLITE_DCHECK(!params.multiplier_exponent_perchannel);
+  }
+}
+
+namespace detail {
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct ValidateTypes {
+  // This generic implementation is for quantized flavors.
+  // kFloatingPoint will be a specialization below.
+  static_assert(!std::is_floating_point<LhsScalar>::value, "");
+  static_assert(!std::is_floating_point<RhsScalar>::value, "");
+  static_assert(!std::is_floating_point<AccumScalar>::value, "");
+  // No requirement on DstScalar --- we might in the future allow it
+  // to be floating point even in a quantized Gemm.
+};
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+struct ValidateTypes<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                     QuantizationFlavor::kFloatingPoint> {
+  static_assert(std::is_floating_point<LhsScalar>::value, "");
+  static_assert(std::is_floating_point<RhsScalar>::value, "");
+  static_assert(std::is_floating_point<AccumScalar>::value, "");
+  static_assert(std::is_floating_point<DstScalar>::value, "");
+};
+
+}  // namespace detail
+
+// Validates overall consistency of all the parameters taken by a Gemm call:
+// the 3 MatrixParams and the GemmParams.
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+void ValidateParams(
+    const MatrixParams<LhsScalar>& lhs_params,
+    const MatrixParams<RhsScalar>& rhs_params,
+    const MatrixParams<DstScalar>& dst_params,
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params) {
+  (void)detail::ValidateTypes<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                              quantization_flavor>();
+  ValidateGemmParams(params);
+}
+
+// Test if the Gemm is degenerate in some way, e.g. nonsensical dimenions.
+template <typename LhsScalar, typename RhsScalar, typename DstScalar>
+bool IsValidGemm(const MatrixParams<LhsScalar>& lhs_params,
+                 const MatrixParams<RhsScalar>& rhs_params,
+                 const MatrixParams<DstScalar>& dst_params) {
+  bool valid = true;
+  valid &= lhs_params.rows >= 1;
+  valid &= lhs_params.cols >= 1;
+  valid &= rhs_params.rows >= 1;
+  valid &= rhs_params.cols >= 1;
+  valid &= dst_params.rows >= 1;
+  valid &= dst_params.cols >= 1;
+  valid &= lhs_params.cols == rhs_params.rows;
+  valid &= rhs_params.cols == dst_params.cols;
+  valid &= lhs_params.rows == lhs_params.rows;
+  return valid;
+}
+
+}  // namespace cpu_backend_gemm
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_PARAMS_H_

tensorflow/lite/kernels/cpu_backend_gemm_ruy.h

@@ -0,0 +1,150 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_RUY_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_RUY_H_
+
+#include "ruy/matrix.h"  // from @ruy
+#include "ruy/mul_params.h"  // from @ruy
+#include "ruy/ruy.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+inline ruy::CachePolicy ToRuyCachePolicy(CachePolicy cache_policy) {
+  switch (cache_policy) {
+    case CachePolicy::kNeverCache:
+      return ruy::CachePolicy::kNeverCache;
+    case CachePolicy::kCacheIfLargeSpeedup:
+      return ruy::CachePolicy::kCacheIfLargeSpeedup;
+    case CachePolicy::kAlwaysCache:
+      return ruy::CachePolicy::kAlwaysCache;
+    default:
+      TFLITE_DCHECK(false);
+      return ruy::CachePolicy::kNeverCache;
+  }
+}
+
+template <typename Scalar, typename DataPointer>
+void MakeRuyMatrix(const MatrixParams<Scalar>& params, DataPointer data_ptr,
+                   ruy::Matrix<Scalar>* dst, bool use_caching = false) {
+  ruy::Order ruy_order = params.order == Order::kColMajor
+                             ? ruy::Order::kColMajor
+                             : ruy::Order::kRowMajor;
+  ruy::MakeSimpleLayout(params.rows, params.cols, ruy_order,
+                        dst->mutable_layout());
+  // Note that ruy::Matrix::data is a ConstCheckingPtr, not a plain pointer.
+  // It does care whether we assign to it a Scalar* or a const Scalar*.
+  dst->set_data(data_ptr);
+  dst->set_zero_point(params.zero_point);
+  if (use_caching) {
+    dst->set_cache_policy(ToRuyCachePolicy(params.cache_policy));
+  }
+}
+
+// Floating-point case.
+template <typename AccumScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+struct MakeRuyMulParamsImpl final {
+  static void Run(
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+      ruy::MulParams<AccumScalar, DstScalar>* ruy_mul_params) {
+    static_assert(quantization_flavor == QuantizationFlavor::kFloatingPoint,
+                  "");
+    ruy_mul_params->set_bias(params.bias);
+    ruy_mul_params->set_clamp_min(params.clamp_min);
+    ruy_mul_params->set_clamp_max(params.clamp_max);
+  }
+};
+
+// Integer-quantized case with destination type narrower than int32
+template <typename DstScalar, QuantizationFlavor quantization_flavor>
+struct MakeRuyMulParamsImpl<std::int32_t, DstScalar, quantization_flavor>
+    final {
+  static void Run(
+      const GemmParams<std::int32_t, DstScalar, quantization_flavor>& params,
+      ruy::MulParams<std::int32_t, DstScalar>* ruy_mul_params) {
+    static_assert(sizeof(DstScalar) < sizeof(std::int32_t), "");
+    if (quantization_flavor ==
+        QuantizationFlavor::kIntegerWithUniformMultiplier) {
+      ruy_mul_params->set_multiplier_fixedpoint(params.multiplier_fixedpoint);
+      ruy_mul_params->set_multiplier_exponent(params.multiplier_exponent);
+    }
+    if (quantization_flavor ==
+        QuantizationFlavor::kIntegerWithPerRowMultiplier) {
+      ruy_mul_params->set_multiplier_fixedpoint_perchannel(
+          params.multiplier_fixedpoint_perchannel);
+      ruy_mul_params->set_multiplier_exponent_perchannel(
+          params.multiplier_exponent_perchannel);
+    }
+    ruy_mul_params->set_bias(params.bias);
+    ruy_mul_params->set_clamp_min(params.clamp_min);
+    ruy_mul_params->set_clamp_max(params.clamp_max);
+  }
+};
+
+// Raw-integer case with destination type int32.
+template <QuantizationFlavor quantization_flavor>
+struct MakeRuyMulParamsImpl<std::int32_t, std::int32_t, quantization_flavor>
+    final {
+  static void Run(
+      const GemmParams<std::int32_t, std::int32_t, quantization_flavor>& params,
+      ruy::MulParams<std::int32_t, std::int32_t>* ruy_mul_params) {
+    ruy_mul_params->set_bias(params.bias);
+  }
+};
+
+template <typename AccumScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+void MakeRuyMulParams(
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+    ruy::MulParams<AccumScalar, DstScalar>* ruy_mul_params) {
+  MakeRuyMulParamsImpl<AccumScalar, DstScalar, quantization_flavor>::Run(
+      params, ruy_mul_params);
+}
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImplUsingRuy {
+  static void Run(
+      const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+      const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+      CpuBackendContext* context) {
+    ruy::Matrix<LhsScalar> ruy_lhs;
+    ruy::Matrix<RhsScalar> ruy_rhs;
+    ruy::Matrix<DstScalar> ruy_dst;
+    MakeRuyMatrix(lhs_params, lhs_data, &ruy_lhs, context->use_caching());
+    MakeRuyMatrix(rhs_params, rhs_data, &ruy_rhs, context->use_caching());
+    MakeRuyMatrix(dst_params, dst_data, &ruy_dst);
+
+    ruy::MulParams<AccumScalar, DstScalar> ruy_mul_params;
+    MakeRuyMulParams(params, &ruy_mul_params);
+
+    ruy::Mul(ruy_lhs, ruy_rhs, ruy_mul_params, context->ruy_context(),
+             &ruy_dst);
+  }
+};
+
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_RUY_H_

tensorflow/lite/kernels/cpu_backend_gemm_test.cc

@@ -0,0 +1,755 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
+
+#include <math.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+#include <algorithm>
+#include <iterator>
+#include <limits>
+#include <random>
+#include <sstream>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "ruy/matrix.h"  // from @ruy
+#include "ruy/reference_mul.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
+
+namespace tflite {
+
+namespace {
+
+using cpu_backend_gemm::Gemm;
+using cpu_backend_gemm::GemmParams;
+using cpu_backend_gemm::MatrixParams;
+using cpu_backend_gemm::QuantizationFlavor;
+
+template <typename Scalar>
+std::string ToString(const std::vector<Scalar>& vector) {
+  std::stringstream s;
+  if (vector.empty()) {
+    s << "{}";
+  } else {
+    s << "{ " << static_cast<double>(vector[0]);
+    for (int i = 1; i < vector.size(); i++) {
+      s << ", " << static_cast<double>(vector[i]);
+    }
+    s << "}";
+  }
+  return s.str();
+}
+
+template <typename Scalar>
+void MakeDeterministicPseudoRandomVector(int size,
+                                         std::vector<Scalar>* vector) {
+  // Intentionally create a new local random_engine in each invocation,
+  // so pseudorandom values don't depend on invocation order.
+  // Otherwise, test results would be affecting by e.g. filtering.
+  std::default_random_engine random_engine;
+  (void)random_engine();
+  // Do not use std::uniform*_distribution: the values that it
+  // generates are implementation-defined.
+  const double random_min = static_cast<double>(random_engine.min());
+  const double random_max = static_cast<double>(random_engine.max());
+  const double result_min =
+      std::is_floating_point<Scalar>::value
+          ? -1.0
+          : std::max(-256., static_cast<double>(
+                                std::numeric_limits<Scalar>::lowest()));
+  const double result_max =
+      std::is_floating_point<Scalar>::value
+          ? 1.0
+          : std::min(256.,
+                     static_cast<double>(std::numeric_limits<Scalar>::max()));
+  const double random_scale =
+      (result_max - result_min) / (random_max - random_min);
+
+  vector->resize(size);
+  for (int i = 0; i < size; i++) {
+    double val = random_scale * (random_engine() - random_min);
+    val = std::max(val,
+                   static_cast<double>(std::numeric_limits<Scalar>::lowest()));
+    val =
+        std::min(val, static_cast<double>(std::numeric_limits<Scalar>::max()));
+    (*vector)[i] = static_cast<Scalar>(val);
+  }
+}
+
+template <typename Scalar>
+void MakeVectorFilledWithConsecutiveInts(int size,
+                                         std::vector<Scalar>* vector) {
+  vector->resize(size);
+  EXPECT_LE(size, std::numeric_limits<Scalar>::max());
+  for (int i = 0; i < size; i++) {
+    (*vector)[i] = static_cast<Scalar>(i + 1);
+  }
+}
+
+template <typename Scalar>
+Scalar Median(const std::vector<Scalar>& vector) {
+  EXPECT_GT(vector.size(), 0);
+  std::vector<Scalar> vector_copy = vector;
+  std::sort(std::begin(vector_copy), std::end(vector_copy));
+  return vector_copy[vector_copy.size() / 2];
+}
+
+template <typename Scalar>
+double MedianAbs(const std::vector<Scalar>& vector) {
+  EXPECT_GT(vector.size(), 0);
+  std::vector<double> vector_abs;
+  vector_abs.resize(vector.size());
+  for (int i = 0; i < vector.size(); i++) {
+    vector_abs[i] = std::abs(static_cast<double>(vector[i]));
+  }
+  std::sort(std::begin(vector_abs), std::end(vector_abs));
+  return vector_abs[vector_abs.size() / 2];
+}
+
+template <typename Scalar>
+void Clamp(const std::vector<Scalar>& src, Scalar clamp_min, Scalar clamp_max,
+           std::vector<Scalar>* dst) {
+  dst->resize(src.size());
+  for (int i = 0; i < src.size(); i++) {
+    (*dst)[i] = std::max(std::min(src[i], clamp_max), clamp_min);
+  }
+}
+
+template <typename AccumScalar, typename DstScalar,
+          QuantizationFlavor quantization_flavor>
+void Clamp(const GemmParams<AccumScalar, DstScalar, quantization_flavor>& src,
+           DstScalar clamp_min, DstScalar clamp_max,
+           GemmParams<AccumScalar, DstScalar, quantization_flavor>* dst) {
+  *dst = src;
+  dst->clamp_min = clamp_min;
+  dst->clamp_max = clamp_max;
+}
+
+struct ErrorStats {
+  int size;
+  double scale_factor;
+  double max_abs_diff;
+  double mean_abs_diff;
+  double abs_mean_diff;
+};
+
+template <typename Scalar>
+void ComputeErrorStats(const std::vector<Scalar>& actual,
+                       const std::vector<Scalar>& expected,
+                       ErrorStats* error_stats) {
+  double max_abs_diff = 0;
+  double sum_abs_diff = 0;
+  double sum_diff = 0;
+  double max_abs_expected = 0;
+  EXPECT_EQ(actual.size(), expected.size());
+  for (int i = 0; i < actual.size(); i++) {
+    double actual_val = static_cast<double>(actual[i]);
+    double expected_val = static_cast<double>(expected[i]);
+    double diff = actual_val - expected_val;
+    max_abs_expected = std::max(max_abs_expected, std::abs(expected_val));
+    sum_diff += diff;
+    sum_abs_diff += std::abs(diff);
+    max_abs_diff = std::max(max_abs_diff, std::abs(diff));
+  }
+  error_stats->scale_factor = max_abs_expected;
+  error_stats->max_abs_diff = max_abs_diff;
+  error_stats->mean_abs_diff = sum_abs_diff / actual.size();
+  error_stats->abs_mean_diff = std::abs(sum_diff / actual.size());
+  error_stats->size = actual.size();
+}
+
+template <typename AccumScalar, typename DstScalar>
+bool CheckErrorStats(const ErrorStats& error_stats, int accumulation_depth) {
+  double tolerated_relative_max_abs_diff = 0;
+  double tolerated_relative_mean_abs_diff = 0;
+  double tolerated_relative_abs_mean_diff = 0;
+
+  double inverse_size = 1. / error_stats.size;
+
+  if (std::is_floating_point<AccumScalar>::value) {
+    // Somewhat naive requirement: the worst case should be epsilons
+    // adding up towards the same direction, on values of same magnitude.
+    tolerated_relative_max_abs_diff =
+        accumulation_depth * std::numeric_limits<DstScalar>::epsilon();
+    // Naive interpretation of the Central Limit Theorem is the rationale
+    // for the sqrt here. We haven't even worked out the correct scale factor,
+    // or how applicable that theorem is here (the random variables being added
+    // might not be mutually independent).
+    tolerated_relative_mean_abs_diff =
+        std::sqrt(static_cast<double>(accumulation_depth)) *
+        std::numeric_limits<DstScalar>::epsilon();
+    // Unbiasing requirement: we require the bias, abs_mean_diff, to be much
+    // smaller than the mean_abs_diff, except when there are very few values.
+    tolerated_relative_abs_mean_diff =
+        tolerated_relative_mean_abs_diff * std::sqrt(inverse_size);
+  } else {
+    // In quantized arithmetic, tolerate minor rounding differences, resulting
+    // in off-by-one errors (tolerated_relative_max_abs_diff = 1), as long
+    // as they are rare (tolerated_relative_mean_abs_diff) and unbiased
+    // (tolerated_relative_abs_mean_diff).
+    tolerated_relative_max_abs_diff = 1;
+    // Naively require mean_abs_diff and abs_mean_diff to converge to zero
+    // as size gets large. We don't know at all how quick that convergence
+    // should be: this is just based on trial-and-error and striking a
+    // compromise between something that works and something that's simple
+    // enough code that doesn't feel too ad-hoc. As above in the float path,
+    // abs_mean_diff is subject to a stricter requirement as it is a bias.
+    tolerated_relative_mean_abs_diff = std::sqrt(inverse_size) * 0.5;
+    tolerated_relative_abs_mean_diff = inverse_size * 2.;
+  }
+
+  double tolerated_max_abs_diff =
+      tolerated_relative_max_abs_diff * error_stats.scale_factor;
+  double tolerated_mean_abs_diff =
+      tolerated_relative_mean_abs_diff * error_stats.scale_factor;
+  double tolerated_abs_mean_diff =
+      tolerated_relative_abs_mean_diff * error_stats.scale_factor;
+
+  EXPECT_LE(error_stats.max_abs_diff, tolerated_max_abs_diff);
+  EXPECT_LE(error_stats.mean_abs_diff, tolerated_mean_abs_diff);
+  EXPECT_LE(error_stats.abs_mean_diff, tolerated_abs_mean_diff);
+
+  return error_stats.max_abs_diff <= tolerated_max_abs_diff &&
+         error_stats.mean_abs_diff <= tolerated_mean_abs_diff &&
+         error_stats.abs_mean_diff <= tolerated_abs_mean_diff;
+}
+
+template <typename AccumScalar, typename DstScalar>
+void CheckErrorForAccumulation(int accumulation_depth,
+                               const std::vector<DstScalar>& actual,
+                               const std::vector<DstScalar>& expected) {
+  ErrorStats error_stats;
+  ComputeErrorStats(actual, expected, &error_stats);
+  bool success =
+      CheckErrorStats<AccumScalar, DstScalar>(error_stats, accumulation_depth);
+  EXPECT_TRUE(success) << "Actual vector\n"
+                       << ToString(actual) << "\ndiffers from expected vector\n"
+                       << ToString(expected) << "\n";
+}
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+void PerformGemmThenCompareResultsThenAgainWithClamping(
+    const MatrixParams<LhsScalar>& lhs_params,
+    const std::vector<LhsScalar>& lhs_data,
+    const MatrixParams<RhsScalar>& rhs_params,
+    const std::vector<RhsScalar>& rhs_data,
+    const MatrixParams<DstScalar>& dst_params, std::vector<DstScalar>* dst_data,
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+    const std::vector<DstScalar>& expected,
+    CpuBackendContext* cpu_backend_context) {
+  const int accumulation_depth = lhs_params.cols;
+  Gemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(), dst_params,
+       dst_data->data(), params, cpu_backend_context);
+  CheckErrorForAccumulation<AccumScalar>(accumulation_depth, *dst_data,
+                                         expected);
+  DstScalar expected_median = Median(expected);
+  std::vector<DstScalar> expected_with_clamp;
+  GemmParams<AccumScalar, DstScalar, quantization_flavor> params_with_clamp;
+  DstScalar clamp_min, clamp_max;
+
+  clamp_min = std::numeric_limits<DstScalar>::lowest();
+  clamp_max = expected_median;
+  Clamp(expected, clamp_min, clamp_max, &expected_with_clamp);
+  Clamp(params, clamp_min, clamp_max, &params_with_clamp);
+  Gemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(), dst_params,
+       dst_data->data(), params_with_clamp, cpu_backend_context);
+  CheckErrorForAccumulation<AccumScalar>(accumulation_depth, *dst_data,
+                                         expected_with_clamp);
+
+  clamp_min = expected_median;
+  clamp_max = std::numeric_limits<DstScalar>::max();
+  Clamp(expected, clamp_min, clamp_max, &expected_with_clamp);
+  Clamp(params, clamp_min, clamp_max, &params_with_clamp);
+  Gemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(), dst_params,
+       dst_data->data(), params_with_clamp, cpu_backend_context);
+  CheckErrorForAccumulation<AccumScalar>(accumulation_depth, *dst_data,
+                                         expected_with_clamp);
+}
+
+// When generating testcases for a quantized GEMM, it's not trivial to
+// pick multiplier exponents: a too low value will result in too many zeros,
+// a too high value will result in too many large clamped values, in both
+// cases testing coverage is harmed. Therefore to ensure good testing coverage
+// we must find a multiplier exponent that's just right.  It would be possible
+// to do so by analysis of the random distribution of values in the result
+// matrix. That however would require some mathematical work that we haven't
+// done so far. Until that is done, the best that we can do is to search for
+// a good exponent value by trial-and-error. This is expensive, as each try
+// requires computing a whole GEMM. This is thus probably a major contribution
+// to the overall latency of this test. To partially mitigate that,
+// we use a bisection to reduce the required number of tries.
+//
+// This function is recursive. The bisect_min and bisect_max arguments
+// are the current bisection bounds. It performs a Gemm with the mid-point,
+// named bisect_mid, as the multiplier exponent. Based on whether the values
+// in the resulting matrix are rather too low or too large in absolute
+// value, it then recurses into the corresponding half of the bisection range.
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+int BisectReasonableMultiplierExponent(
+    int bisect_min, int bisect_max, const MatrixParams<LhsScalar>& lhs_params,
+    const std::vector<LhsScalar>& lhs_data,
+    const MatrixParams<RhsScalar>& rhs_params,
+    const std::vector<RhsScalar>& rhs_data,
+    const MatrixParams<DstScalar>& dst_params, std::vector<DstScalar>* dst_data,
+    const GemmParams<AccumScalar, DstScalar>& params,
+    CpuBackendContext* cpu_backend_context) {
+  if (bisect_min == bisect_max) {
+    return bisect_min;
+  }
+  // Compute the midpoint as the floor of the average of bisect_min and
+  // bisect_max. As C++ integer division is rounding towards zero and our values
+  // may be of any sign, it is not trivial to implement this using only integer
+  // arithmetic.
+  int bisect_mid =
+      static_cast<int>(std::floor(0.5 * (bisect_min + bisect_max)));
+  GemmParams<AccumScalar, DstScalar> params_copy(params);
+  params_copy.multiplier_exponent = bisect_mid;
+  double clamp_abs = std::max(std::abs(static_cast<double>(params.clamp_min)),
+                              std::abs(static_cast<double>(params.clamp_max)));
+  Gemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(), dst_params,
+       dst_data->data(), params_copy, cpu_backend_context);
+  double median_abs = MedianAbs(*dst_data);
+  if (median_abs < 0.25 * clamp_abs) {
+    return BisectReasonableMultiplierExponent(
+        bisect_mid + 1, bisect_max, lhs_params, lhs_data, rhs_params, rhs_data,
+        dst_params, dst_data, params_copy, cpu_backend_context);
+  } else {
+    return BisectReasonableMultiplierExponent(
+        bisect_min, bisect_mid, lhs_params, lhs_data, rhs_params, rhs_data,
+        dst_params, dst_data, params_copy, cpu_backend_context);
+  }
+}
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+void ReferenceGemm(
+    const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+    const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+    const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+    const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+    CpuBackendContext* context) {
+  ruy::Matrix<LhsScalar> ruy_lhs;
+  ruy::Matrix<RhsScalar> ruy_rhs;
+  ruy::Matrix<DstScalar> ruy_dst;
+  cpu_backend_gemm::detail::MakeRuyMatrix(lhs_params, lhs_data, &ruy_lhs);
+  cpu_backend_gemm::detail::MakeRuyMatrix(rhs_params, rhs_data, &ruy_rhs);
+  cpu_backend_gemm::detail::MakeRuyMatrix(dst_params, dst_data, &ruy_dst);
+
+  ruy::MulParams<AccumScalar, DstScalar> ruy_mul_params;
+  cpu_backend_gemm::detail::MakeRuyMulParams(params, &ruy_mul_params);
+
+  ruy::ReferenceMul(ruy_lhs, ruy_rhs, ruy_mul_params, &ruy_dst);
+}
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+void TestSomeGemm(int rows, int depth, int cols,
+                  const std::vector<DstScalar>& golden) {
+  CpuBackendContext cpu_backend_context;
+  std::default_random_engine random_engine;
+  cpu_backend_context.SetMaxNumThreads(1 + (random_engine() % 8));
+  bool use_caching = static_cast<bool>(random_engine() % 2);
+  cpu_backend_context.SetUseCaching(use_caching);
+  const bool use_golden = !golden.empty();
+
+  std::vector<LhsScalar> lhs_data;
+  std::vector<RhsScalar> rhs_data;
+  std::vector<AccumScalar> bias_data;
+  std::vector<DstScalar> dst_data;
+  if (use_golden) {
+    MakeVectorFilledWithConsecutiveInts(rows * depth, &lhs_data);
+    MakeVectorFilledWithConsecutiveInts(depth * cols, &rhs_data);
+    MakeVectorFilledWithConsecutiveInts(rows, &bias_data);
+  } else {
+    MakeDeterministicPseudoRandomVector(rows * depth, &lhs_data);
+    MakeDeterministicPseudoRandomVector(depth * cols, &rhs_data);
+    MakeDeterministicPseudoRandomVector(rows, &bias_data);
+  }
+  MakeDeterministicPseudoRandomVector(rows * cols, &dst_data);
+
+  auto random_order = [&]() {
+    return random_engine() % 2 ? cpu_backend_gemm::Order::kRowMajor
+                               : cpu_backend_gemm::Order::kColMajor;
+  };
+  MatrixParams<LhsScalar> lhs_params;
+  lhs_params.order =
+      use_golden ? cpu_backend_gemm::Order::kRowMajor : random_order();
+  lhs_params.rows = rows;
+  lhs_params.cols = depth;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<LhsScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    lhs_params.zero_point = 1;
+    if (!use_golden) {
+      lhs_params.zero_point += random_engine() % 8;
+    }
+  }
+
+  MatrixParams<RhsScalar> rhs_params;
+  rhs_params.order =
+      use_golden ? cpu_backend_gemm::Order::kColMajor : random_order();
+  rhs_params.rows = depth;
+  rhs_params.cols = cols;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<RhsScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    rhs_params.zero_point = 1;
+    if (!use_golden) {
+      rhs_params.zero_point += random_engine() % 8;
+    }
+  }
+
+  MatrixParams<DstScalar> dst_params;
+  dst_params.order =
+      use_golden ? cpu_backend_gemm::Order::kColMajor : random_order();
+  dst_params.rows = rows;
+  dst_params.cols = cols;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<DstScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    dst_params.zero_point = 1;
+    if (!use_golden) {
+      dst_params.zero_point += random_engine() % 8;
+    }
+  }
+
+  GemmParams<AccumScalar, DstScalar> params;
+  if (use_golden || (random_engine() % 2)) {
+    // cpu_backend_gemm supports bias=null only in the float path. Test that
+    // in 50% of float testcases.
+    params.bias = bias_data.data();
+  }
+  static constexpr std::int32_t kMultiplierFixedpointMin = 1234567890;
+  static constexpr std::int32_t kMultiplierFixedpointMax = 1987654321;
+  if (!std::is_floating_point<AccumScalar>::value) {
+    // some large int32 value. Not being a multiple of a large
+    // power of two helps testing rounding behavior.
+    params.multiplier_fixedpoint = kMultiplierFixedpointMin;
+    // Now find a suitable value for multiplier_exponent.
+    // It needs to be low enough for a substantial amount of dst values
+    // to avoid getting clamped.
+    int bisect_min = -8 * static_cast<int>(sizeof(AccumScalar));
+    // We don't increase test coverage by using positive multipliers,
+    // and using very large positive multipliers may at the moment
+    // result in overflow in some paths.
+    // TODO(benoitjacob): fix that.
+    int bisect_max = 0;
+    params.multiplier_exponent = BisectReasonableMultiplierExponent(
+        bisect_min, bisect_max, lhs_params, lhs_data, rhs_params, rhs_data,
+        dst_params, &dst_data, params, &cpu_backend_context);
+  }
+
+  std::vector<DstScalar> expected;
+  if (use_golden) {
+    EXPECT_EQ(golden.size(), dst_data.size());
+    expected = golden;
+  } else {
+    expected.resize(dst_data.size());
+    ReferenceGemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(),
+                  dst_params, expected.data(), params, &cpu_backend_context);
+  }
+
+  PerformGemmThenCompareResultsThenAgainWithClamping(
+      lhs_params, lhs_data, rhs_params, rhs_data, dst_params, &dst_data, params,
+      expected, &cpu_backend_context);
+
+  if (!use_golden && !std::is_floating_point<AccumScalar>::value) {
+    // Try with per-channel quantized multipliers.
+    std::vector<AccumScalar> multiplier_fixedpoint_perchannel(rows);
+    std::vector<int> multiplier_exponent_perchannel(rows);
+    for (int i = 0; i < rows; i++) {
+      multiplier_fixedpoint_perchannel[i] =
+          kMultiplierFixedpointMin +
+          (random_engine() %
+           (kMultiplierFixedpointMax + 1 - kMultiplierFixedpointMin));
+      const int exponent_min = params.multiplier_exponent - 2;
+      const int exponent_max = params.multiplier_exponent + 2;
+      multiplier_exponent_perchannel[i] =
+          exponent_min + (random_engine() % (exponent_max + 1 - exponent_min));
+    }
+    static constexpr QuantizationFlavor perchannel_flavor =
+        std::is_floating_point<AccumScalar>::value
+            ? QuantizationFlavor::kFloatingPoint
+            : QuantizationFlavor::kIntegerWithPerRowMultiplier;
+    GemmParams<AccumScalar, DstScalar, perchannel_flavor> params_perchannel;
+    params_perchannel.bias = params.bias;
+    params_perchannel.clamp_min = params.clamp_min;
+    params_perchannel.clamp_max = params.clamp_max;
+    params_perchannel.multiplier_fixedpoint_perchannel =
+        multiplier_fixedpoint_perchannel.data();
+    params_perchannel.multiplier_exponent_perchannel =
+        multiplier_exponent_perchannel.data();
+    ReferenceGemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(),
+                  dst_params, expected.data(), params_perchannel,
+                  &cpu_backend_context);
+    PerformGemmThenCompareResultsThenAgainWithClamping(
+        lhs_params, lhs_data, rhs_params, rhs_data, dst_params, &dst_data,
+        params_perchannel, expected, &cpu_backend_context);
+  }
+}
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar>
+void TestMaybeValidGemm(int lhs_rows, int lhs_cols, int rhs_rows, int rhs_cols,
+                        int dst_rows, int dst_cols) {
+  CpuBackendContext cpu_backend_context;
+  std::default_random_engine random_engine;
+  cpu_backend_context.SetMaxNumThreads(1 + (random_engine() % 8));
+  bool use_caching = static_cast<bool>(random_engine() % 2);
+  cpu_backend_context.SetUseCaching(use_caching);
+
+  std::vector<LhsScalar> lhs_data;
+  std::vector<RhsScalar> rhs_data;
+  std::vector<AccumScalar> bias_data;
+  std::vector<DstScalar> dst_data;
+  MakeDeterministicPseudoRandomVector(lhs_rows * lhs_cols, &lhs_data);
+  MakeDeterministicPseudoRandomVector(rhs_rows * rhs_cols, &rhs_data);
+  MakeDeterministicPseudoRandomVector(dst_rows, &bias_data);
+  MakeDeterministicPseudoRandomVector(dst_rows * dst_cols, &dst_data);
+
+  MatrixParams<LhsScalar> lhs_params;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.rows = lhs_rows;
+  lhs_params.cols = lhs_cols;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<LhsScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    lhs_params.zero_point = 1;
+  }
+
+  MatrixParams<RhsScalar> rhs_params;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.rows = rhs_rows;
+  rhs_params.cols = rhs_cols;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<RhsScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    rhs_params.zero_point = 1;
+  }
+
+  MatrixParams<DstScalar> dst_params;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.rows = dst_rows;
+  dst_params.cols = dst_cols;
+  // 16x8 quant only supports RUY path. For Ruy 16x8 gemm, it restricts
+  // zero_point as 0 because int16 might cause overflow in acuum int32.
+  // https://github.com/google/ruy/blob/master/ruy/validate.h#L53-L57
+  if (!std::is_floating_point<DstScalar>::value &&
+      (!std::is_same<LhsScalar, int8_t>::value &&
+       !std::is_same<RhsScalar, int16_t>::value)) {
+    dst_params.zero_point = 1;
+  }
+
+  GemmParams<AccumScalar, DstScalar> params;
+  params.bias = bias_data.data();
+  static constexpr std::int32_t kMultiplierFixedpointMin = 1234567890;
+  if (!std::is_floating_point<AccumScalar>::value) {
+    // some large int32 value. Not being a multiple of a large
+    // power of two helps testing rounding behavior.
+    params.multiplier_fixedpoint = kMultiplierFixedpointMin;
+    // Now find a suitable value for multiplier_exponent.
+    // It needs to be low enough for a substantial amount of dst values
+    // to avoid getting clamped.
+    int bisect_min = -8 * static_cast<int>(sizeof(AccumScalar));
+    // We don't increase test coverage by using positive multipliers,
+    // and using very large positive multipliers may at the moment
+    // result in overflow in some paths.
+    int bisect_max = 0;
+    params.multiplier_exponent = BisectReasonableMultiplierExponent(
+        bisect_min, bisect_max, lhs_params, lhs_data, rhs_params, rhs_data,
+        dst_params, &dst_data, params, &cpu_backend_context);
+  }
+  Gemm(lhs_params, lhs_data.data(), rhs_params, rhs_data.data(), dst_params,
+       dst_data.data(), params, &cpu_backend_context);
+}
+
+TEST(CpuBackendGemmSimpleTestAgainstGolden, Float) {
+  TestSomeGemm<float, float, float, float>(2, 3, 4,
+                                           {15, 34, 33, 79, 51, 124, 69, 169});
+}
+
+TEST(CpuBackendGemmSimpleTestAgainstGolden, Uint8) {
+  TestSomeGemm<std::uint8_t, std::uint8_t, std::int32_t, std::uint8_t>(
+      5, 2, 3, {2, 4, 6, 7, 9, 3, 10, 16, 22, 29, 4, 15, 26, 37, 48});
+}
+
+TEST(CpuBackendGemmSimpleTestAgainstGolden, Int8) {
+  TestSomeGemm<std::int8_t, std::int8_t, std::int32_t, std::int8_t>(
+      2, 6, 3, {13, 32, 31, 81, 50, 127});
+}
+
+TEST(CpuBackendGemmInvalidGemmTest, Float) {
+  // A standard Gemm operation.
+  TestMaybeValidGemm<float, float, float, float>(2, 3, 3, 4, 2, 4);
+  // An invalid Gemm that will abort in debug mode.
+#if !defined(TARGET_IPHONE_SIMULATOR) && !defined(TARGET_OS_IPHONE)
+  ASSERT_DEBUG_DEATH(
+      (TestMaybeValidGemm<float, float, float, float>(2, 3, 3, 0, 2, 4)), "");
+  ASSERT_DEBUG_DEATH(
+      (TestMaybeValidGemm<float, float, float, float>(2, 3, 9, 4, 2, 4)), "");
+#endif
+}
+
+TEST(CpuBackendGemmSimpleTestAgainstGolden, Int8Int16) {
+  TestSomeGemm<std::int8_t, std::int8_t, std::int32_t, std::int16_t>(
+      3, 5, 4, {19, 48, 77, 48, 149, 250, 76, 249, 422, 105, 350, 595});
+}
+
+template <typename tLhsScalar, typename tRhsScalar, typename tAccumScalar,
+          typename tDstScalar>
+struct TypesTuple {
+  using LhsScalar = tLhsScalar;
+  using RhsScalar = tRhsScalar;
+  using AccumScalar = tAccumScalar;
+  using DstScalar = tDstScalar;
+};
+
+template <typename TypesTupleType>
+void TestRandomGemms(const std::vector<std::tuple<int, int, int>>& shapes) {
+  using LhsScalar = typename TypesTupleType::LhsScalar;
+  using RhsScalar = typename TypesTupleType::RhsScalar;
+  using AccumScalar = typename TypesTupleType::AccumScalar;
+  using DstScalar = typename TypesTupleType::DstScalar;
+  for (const auto& shape : shapes) {
+    int rows = std::get<0>(shape);
+    int depth = std::get<1>(shape);
+    int cols = std::get<2>(shape);
+    TestSomeGemm<LhsScalar, RhsScalar, AccumScalar, DstScalar>(rows, depth,
+                                                               cols, {});
+  }
+}
+
+template <typename TypesTupleType>
+class CpuBackendGemmTest : public testing::Test {};
+
+TYPED_TEST_SUITE_P(CpuBackendGemmTest);
+
+typedef ::testing::Types<
+    TypesTuple<float, float, float, float>,
+    TypesTuple<std::uint8_t, std::uint8_t, std::int32_t, std::uint8_t>,
+    TypesTuple<std::int8_t, std::int8_t, std::int32_t, std::int8_t>,
+    TypesTuple<std::int8_t, std::int8_t, std::int32_t, std::int16_t>,
+    TypesTuple<std::int8_t, std::int16_t, std::int32_t, std::int16_t>,
+    TypesTuple<std::uint8_t, std::uint8_t, std::int32_t, std::int8_t>>
+    CpuBackendGemmTestInstantiations;
+
+TYPED_TEST_SUITE(CpuBackendGemmTest, CpuBackendGemmTestInstantiations);
+
+TYPED_TEST(CpuBackendGemmTest, Square) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 50; size++) {
+    shapes.push_back(std::make_tuple(size, size, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, SquarePowerOfTwo) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 64; size <= 128; size *= 2) {
+    shapes.push_back(std::make_tuple(size, size, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, MatrixTimesVector) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 200; size++) {
+    shapes.push_back(std::make_tuple(size, size, 1));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, VectorTimesMatrix) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 200; size++) {
+    shapes.push_back(std::make_tuple(1, size, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, MatrixTimesNarrow) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 50; size++) {
+    shapes.push_back(std::make_tuple(size, size, 2));
+    shapes.push_back(std::make_tuple(size, size, 3));
+    shapes.push_back(std::make_tuple(size, size, 4));
+    shapes.push_back(std::make_tuple(size, size, 8));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, Rectangular) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 50; size++) {
+    shapes.push_back(std::make_tuple(size, size + 5, size + 1));
+    shapes.push_back(std::make_tuple(size + 10, size + 2, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, HighlyRectangular) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size <= 1000; size *= 10) {
+    shapes.push_back(std::make_tuple(size, 10, 10));
+    shapes.push_back(std::make_tuple(10, size, 10));
+    shapes.push_back(std::make_tuple(10, 10, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, InnerProduct) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 200; size++) {
+    shapes.push_back(std::make_tuple(1, size, 1));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+TYPED_TEST(CpuBackendGemmTest, OuterProduct) {
+  std::vector<std::tuple<int, int, int>> shapes;
+  for (int size = 1; size < 100; size++) {
+    shapes.push_back(std::make_tuple(size, 1, size));
+  }
+  TestRandomGemms<TypeParam>(shapes);
+}
+
+}  // namespace
+
+}  // namespace tflite

tensorflow/lite/kernels/cpu_backend_gemm_x86.h

@@ -0,0 +1,106 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_X86_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_X86_H_
+
+// If TFLITE_WITH_RUY is set, Ruy is the only GEMM option. In this header
+// we select either Ruy or an alternative based on the SIMD extentions
+// available on the given x86 platform.
+#ifndef TFLITE_WITH_RUY
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_eigen.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace cpu_backend_gemm {
+namespace detail {
+
+template <typename LhsScalar, typename RhsScalar, typename AccumScalar,
+          typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImplX86 {
+  static void Run(
+      const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data,
+      const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data,
+      const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data,
+      const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params,
+      CpuBackendContext* context) {
+    // TODO(b/168923364) Ruy is preferred on x86, but check if the deprecated
+    // path is enabled.
+    if (context->PreferGemmlowpOnX86()) {
+      // Dispatch to gemmlowp.
+      detail::GemmImplUsingGemmlowp<
+          LhsScalar, RhsScalar, AccumScalar, DstScalar,
+          quantization_flavor>::Run(lhs_params, lhs_data, rhs_params, rhs_data,
+                                    dst_params, dst_data, params, context);
+
+      return;
+    }
+    // Run-time dispatch to Ruy for platforms with AVX or above.
+    detail::GemmImplUsingRuy<LhsScalar, RhsScalar, AccumScalar, DstScalar,
+                             quantization_flavor>::Run(lhs_params, lhs_data,
+                                                       rhs_params, rhs_data,
+                                                       dst_params, dst_data,
+                                                       params, context);
+  }
+};
+
+// For float, defer to eigen for now.
+template <>
+struct GemmImplX86<float, float, float, float,
+                   QuantizationFlavor::kFloatingPoint> {
+  static void Run(const MatrixParams<float>& lhs_params, const float* lhs_data,
+                  const MatrixParams<float>& rhs_params, const float* rhs_data,
+                  const MatrixParams<float>& dst_params, float* dst_data,
+                  const GemmParams<float, float,
+                                   QuantizationFlavor::kFloatingPoint>& params,
+                  CpuBackendContext* context) {
+    GemmImplUsingEigen::Run(lhs_params, lhs_data, rhs_params, rhs_data,
+                            dst_params, dst_data, params, context);
+  }
+};
+
+// gemmlowp requires NEON for certain quantization cases. See note in
+// cpu_backend_gemm.h
+#if !defined(GEMMLOWP_NEON)
+template <typename SrcScalar, QuantizationFlavor quantization_flavor>
+struct GemmImplX86<SrcScalar, SrcScalar, std::int32_t, std::int8_t,
+                   quantization_flavor>
+    : detail::GemmImplUsingRuy<SrcScalar, SrcScalar, std::int32_t, std::int8_t,
+                               quantization_flavor> {};
+
+template <typename DstScalar, QuantizationFlavor quantization_flavor>
+struct GemmImplX86<std::int8_t, std::int8_t, std::int32_t, DstScalar,
+                   quantization_flavor>
+    : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t,
+                               DstScalar, quantization_flavor> {};
+
+template <QuantizationFlavor quantization_flavor>
+struct GemmImplX86<std::int8_t, std::int8_t, std::int32_t, std::int8_t,
+                   quantization_flavor>
+    : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t,
+                               std::int8_t, quantization_flavor> {};
+#endif  // not GEMMLOWP_NEON
+}  // namespace detail
+}  // namespace cpu_backend_gemm
+}  // namespace tflite
+
+#endif  // not TFLITE_WITH_RUY
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_X86_H_

tensorflow/lite/kernels/cpu_backend_threadpool.h

@@ -0,0 +1,61 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_
+#define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+#ifdef TFLITE_WITH_RUY
+#include "ruy/context.h"  // from @ruy
+#include "ruy/thread_pool.h"  // from @ruy
+#else
+#include "public/gemmlowp.h"
+#endif
+
+namespace tflite {
+namespace cpu_backend_threadpool {
+
+#ifdef TFLITE_WITH_RUY
+
+using Task = ruy::Task;
+
+template <typename TaskType>
+void Execute(int tasks_count, TaskType* tasks,
+             CpuBackendContext* cpu_backend_context) {
+  TFLITE_DCHECK_LE(tasks_count, cpu_backend_context->max_num_threads());
+  cpu_backend_context->ruy_context()->mutable_thread_pool()->Execute(
+      tasks_count, tasks);
+}
+
+#else  // not TFLITE_WITH_RUY
+
+using Task = gemmlowp::Task;
+
+template <typename TaskType>
+void Execute(int tasks_count, TaskType* tasks,
+             CpuBackendContext* cpu_backend_context) {
+  TFLITE_DCHECK_LE(tasks_count, cpu_backend_context->max_num_threads());
+  cpu_backend_context->gemmlowp_context()->workers_pool()->Execute(tasks_count,
+                                                                   tasks);
+}
+
+#endif
+
+}  // namespace cpu_backend_threadpool
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_THREADPOOL_H_

tensorflow/lite/kernels/cpu_backend_threadpool_test.cc

@@ -0,0 +1,94 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+
+namespace tflite {
+
+namespace {
+
+class TestGenerateArrayOfIncrementingIntsTask
+    : public cpu_backend_threadpool::Task {
+ public:
+  TestGenerateArrayOfIncrementingIntsTask(int* buffer, int start, int end)
+      : buffer_(buffer), start_(start), end_(end) {}
+
+  void Run() override {
+    for (int i = start_; i < end_; i++) {
+      buffer_[i] = i;
+    }
+  }
+
+ private:
+  int* buffer_;
+  int start_;
+  int end_;
+};
+
+void TestGenerateArrayOfIncrementingInts(int num_threads, int size) {
+  // The buffer that our threads will write to.
+  std::vector<int> buffer(size);
+
+  // The tasks that our threads will run.
+  std::vector<TestGenerateArrayOfIncrementingIntsTask> tasks;
+
+  // Create task objects.
+  int rough_size_per_thread = size / num_threads;
+  int start = 0;
+  for (int thread = 0; thread < num_threads; thread++) {
+    int end = start + rough_size_per_thread;
+    if (thread == num_threads - 1) {
+      end = size;
+    }
+    tasks.emplace_back(buffer.data(), start, end);
+    start = end;
+  }
+  ASSERT_EQ(num_threads, tasks.size());
+
+  CpuBackendContext context;
+  // This SetMaxNumThreads is only to satisfy an assertion in Execute.
+  // What actually determines the number of threads used is the parameter
+  // passed to Execute, since Execute does 1:1 mapping of tasks to threads.
+  context.SetMaxNumThreads(num_threads);
+
+  // Execute tasks on the threadpool.
+  cpu_backend_threadpool::Execute(tasks.size(), tasks.data(), &context);
+
+  // Check contents of the generated buffer.
+  for (int i = 0; i < size; i++) {
+    ASSERT_EQ(buffer[i], i);
+  }
+}
+
+TEST(CpuBackendThreadpoolTest, OneThreadSize100) {
+  TestGenerateArrayOfIncrementingInts(1, 100);
+}
+
+TEST(CpuBackendThreadpoolTest, ThreeThreadsSize1000000) {
+  TestGenerateArrayOfIncrementingInts(3, 1000000);
+}
+
+TEST(CpuBackendThreadpoolTest, TenThreadsSize1234567) {
+  TestGenerateArrayOfIncrementingInts(10, 1234567);
+}
+
+}  // namespace
+
+}  // namespace tflite

tensorflow/lite/kernels/ctc/BUILD

@@ -0,0 +1,80 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+# ctc support classes imported directly from TensorFlow.
+cc_library(
+    name = "ctc_utils",
+    hdrs = [
+        "ctc_beam_entry.h",
+        "ctc_beam_scorer.h",
+        "ctc_beam_search.h",
+        "ctc_decoder.h",
+        "ctc_loss_util.h",
+    ],
+    deps = [
+        ":top_n",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+# top_n support classes imported directly from TensorFlow.
+cc_library(
+    name = "top_n",
+    hdrs = [
+        "top_n.h",
+    ],
+    deps = [
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_library(
+    name = "ctc_beam_search_decoder_op",
+    srcs = [
+        "ctc_beam_search_decoder.cc",
+    ],
+    # Suppress warnings that are introduced by Eigen Tensor.
+    copts = tflite_copts() + [
+        "-Wno-error=reorder",
+    ] + select({
+        "//tensorflow:ios": ["-Wno-error=invalid-partial-specialization"],
+        "//conditions:default": [
+        ],
+    }),
+    deps = [
+        ":ctc_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels:op_macros",
+        "//tensorflow/lite/kernels/internal:optimized_base",
+        "//tensorflow/lite/kernels/internal:tensor",
+        "@flatbuffers",
+    ],
+)
+
+cc_test(
+    name = "ctc_beam_search_decoder_test",
+    size = "small",
+    srcs = ["ctc_beam_search_decoder_test.cc"],
+    tags = ["tflite_not_portable_ios"],
+    deps = [
+        ":ctc_beam_search_decoder_op",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)

tensorflow/lite/kernels/ctc/ctc_beam_entry.h

@@ -0,0 +1,150 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Copied from tensorflow/core/util/ctc/ctc_beam_entry.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_ENTRY_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_ENTRY_H_
+
+#include <algorithm>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/kernels/ctc/ctc_loss_util.h"
+
+namespace tflite {
+namespace custom {
+namespace ctc {
+
+// The ctc_beam_search namespace holds several classes meant to be accessed only
+// in case of extending the CTCBeamSearch decoder to allow custom scoring
+// functions.
+//
+// BeamEntry is exposed through template arguments BeamScorer and BeamComparer
+// of CTCBeamSearch (ctc_beam_search.h).
+namespace ctc_beam_search {
+
+struct EmptyBeamState {};
+
+struct BeamProbability {
+  BeamProbability() : total(kLogZero), blank(kLogZero), label(kLogZero) {}
+  void Reset() {
+    total = kLogZero;
+    blank = kLogZero;
+    label = kLogZero;
+  }
+  float total;
+  float blank;
+  float label;
+};
+
+template <class CTCBeamState>
+class BeamRoot;
+
+template <class CTCBeamState = EmptyBeamState>
+struct BeamEntry {
+  // BeamRoot<CTCBeamState>::AddEntry() serves as the factory method.
+  friend BeamEntry<CTCBeamState>* BeamRoot<CTCBeamState>::AddEntry(
+      BeamEntry<CTCBeamState>* p, int l);
+  inline bool Active() const { return newp.total != kLogZero; }
+  // Return the child at the given index, or construct a new one in-place if
+  // none was found.
+  BeamEntry& GetChild(int ind) {
+    auto entry = children.emplace(ind, nullptr);
+    auto& child_entry = entry.first->second;
+    // If this is a new child, populate the BeamEntry<CTCBeamState>*.
+    if (entry.second) {
+      child_entry = beam_root->AddEntry(this, ind);
+    }
+    return *child_entry;
+  }
+  std::vector<int> LabelSeq(bool merge_repeated) const {
+    std::vector<int> labels;
+    int prev_label = -1;
+    const BeamEntry* c = this;
+    while (c->parent != nullptr) {  // Checking c->parent to skip root leaf.
+      if (!merge_repeated || c->label != prev_label) {
+        labels.push_back(c->label);
+      }
+      prev_label = c->label;
+      c = c->parent;
+    }
+    std::reverse(labels.begin(), labels.end());
+    return labels;
+  }
+
+  BeamEntry<CTCBeamState>* parent;
+  int label;
+  // All instances of child BeamEntry are owned by *beam_root.
+  std::unordered_map<int, BeamEntry<CTCBeamState>*> children;
+  BeamProbability oldp;
+  BeamProbability newp;
+  CTCBeamState state;
+
+ private:
+  // Constructor giving parent, label, and the beam_root.
+  // The object pointed to by p cannot be copied and should not be moved,
+  // otherwise parent will become invalid.
+  // This private constructor is only called through the factory method
+  // BeamRoot<CTCBeamState>::AddEntry().
+  BeamEntry(BeamEntry* p, int l, BeamRoot<CTCBeamState>* beam_root)
+      : parent(p), label(l), beam_root(beam_root) {}
+  BeamRoot<CTCBeamState>* beam_root;
+
+  BeamEntry(const BeamEntry&) = delete;
+  void operator=(const BeamEntry&) = delete;
+};
+
+// This class owns all instances of BeamEntry.  This is used to avoid recursive
+// destructor call during destruction.
+template <class CTCBeamState = EmptyBeamState>
+class BeamRoot {
+ public:
+  BeamRoot(BeamEntry<CTCBeamState>* p, int l) { root_entry_ = AddEntry(p, l); }
+  BeamRoot(const BeamRoot&) = delete;
+  BeamRoot& operator=(const BeamRoot&) = delete;
+
+  BeamEntry<CTCBeamState>* AddEntry(BeamEntry<CTCBeamState>* p, int l) {
+    auto* new_entry = new BeamEntry<CTCBeamState>(p, l, this);
+    beam_entries_.emplace_back(new_entry);
+    return new_entry;
+  }
+  BeamEntry<CTCBeamState>* RootEntry() const { return root_entry_; }
+
+ private:
+  BeamEntry<CTCBeamState>* root_entry_ = nullptr;
+  std::vector<std::unique_ptr<BeamEntry<CTCBeamState>>> beam_entries_;
+};
+
+// BeamComparer is the default beam comparer provided in CTCBeamSearch.
+template <class CTCBeamState = EmptyBeamState>
+class BeamComparer {
+ public:
+  virtual ~BeamComparer() {}
+  virtual bool inline operator()(const BeamEntry<CTCBeamState>* a,
+                                 const BeamEntry<CTCBeamState>* b) const {
+    return a->newp.total > b->newp.total;
+  }
+};
+
+}  // namespace ctc_beam_search
+
+}  // namespace ctc
+}  // namespace custom
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_ENTRY_H_

tensorflow/lite/kernels/ctc/ctc_beam_scorer.h

@@ -0,0 +1,79 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Collection of scoring classes that can be extended and provided to the
+// CTCBeamSearchDecoder to incorporate additional scoring logic (such as a
+// language model).
+//
+// To build a custom scorer extend and implement the pure virtual methods from
+// BeamScorerInterface. The default CTC decoding behavior is implemented
+// through BaseBeamScorer.
+
+// Copied from tensorflow/core/util/ctc/ctc_beam_scorer.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SCORER_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SCORER_H_
+
+#include "tensorflow/lite/kernels/ctc/ctc_beam_entry.h"
+
+namespace tflite {
+namespace custom {
+namespace ctc {
+
+// Base implementation of a beam scorer used by default by the decoder that can
+// be subclassed and provided as an argument to CTCBeamSearchDecoder, if complex
+// scoring is required. Its main purpose is to provide a thin layer for
+// integrating language model scoring easily.
+template <typename CTCBeamState>
+class BaseBeamScorer {
+ public:
+  virtual ~BaseBeamScorer() {}
+  // State initialization.
+  virtual void InitializeState(CTCBeamState* root) const {}
+  // ExpandState is called when expanding a beam to one of its children.
+  // Called at most once per child beam. In the simplest case, no state
+  // expansion is done.
+  virtual void ExpandState(const CTCBeamState& from_state, int from_label,
+                           CTCBeamState* to_state, int to_label) const {}
+  // ExpandStateEnd is called after decoding has finished. Its purpose is to
+  // allow a final scoring of the beam in its current state, before resorting
+  // and retrieving the TopN requested candidates. Called at most once per beam.
+  virtual void ExpandStateEnd(CTCBeamState* state) const {}
+  // GetStateExpansionScore should be an inexpensive method to retrieve the
+  // (cached) expansion score computed within ExpandState. The score is
+  // multiplied (log-addition) with the input score at the current step from
+  // the network.
+  //
+  // The score returned should be a log-probability. In the simplest case, as
+  // there's no state expansion logic, the expansion score is zero.
+  virtual float GetStateExpansionScore(const CTCBeamState& state,
+                                       float previous_score) const {
+    return previous_score;
+  }
+  // GetStateEndExpansionScore should be an inexpensive method to retrieve the
+  // (cached) expansion score computed within ExpandStateEnd. The score is
+  // multiplied (log-addition) with the final probability of the beam.
+  //
+  // The score returned should be a log-probability.
+  virtual float GetStateEndExpansionScore(const CTCBeamState& state) const {
+    return 0;
+  }
+};
+
+}  // namespace ctc
+}  // namespace custom
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SCORER_H_

tensorflow/lite/kernels/ctc/ctc_beam_search.h

@@ -0,0 +1,432 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Copied from tensorflow/core/util/ctc/ctc_beam_search.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SEARCH_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SEARCH_H_
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/kernels/ctc/ctc_beam_entry.h"
+#include "tensorflow/lite/kernels/ctc/ctc_beam_scorer.h"
+#include "tensorflow/lite/kernels/ctc/ctc_decoder.h"
+#include "tensorflow/lite/kernels/ctc/ctc_loss_util.h"
+#include "tensorflow/lite/kernels/ctc/top_n.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace custom {
+namespace ctc {
+
+template <typename CTCBeamState = ctc_beam_search::EmptyBeamState,
+          typename CTCBeamComparer =
+              ctc_beam_search::BeamComparer<CTCBeamState>>
+class CTCBeamSearchDecoder : public CTCDecoder {
+  // Beam Search
+  //
+  // Example (GravesTh Fig. 7.5):
+  //         a    -
+  //  P = [ 0.3  0.7 ]  t = 0
+  //      [ 0.4  0.6 ]  t = 1
+  //
+  // Then P(l = -) = P(--) = 0.7 * 0.6 = 0.42
+  //      P(l = a) = P(a-) + P(aa) + P(-a) = 0.3*0.4 + ... = 0.58
+  //
+  // In this case, Best Path decoding is suboptimal.
+  //
+  // For Beam Search, we use the following main recurrence relations:
+  //
+  // Relation 1:
+  // ---------------------------------------------------------- Eq. 1
+  //      P(l=abcd @ t=7) = P(l=abc  @ t=6) * P(d @ 7)
+  //                      + P(l=abcd @ t=6) * (P(d @ 7) + P(- @ 7))
+  // where P(l=? @ t=7), ? = a, ab, abc, abcd are all stored and
+  // updated recursively in the beam entry.
+  //
+  // Relation 2:
+  // ---------------------------------------------------------- Eq. 2
+  //      P(l=abc? @ t=3) = P(l=abc @ t=2) * P(? @ 3)
+  // for ? in a, b, d, ..., (not including c or the blank index),
+  // and the recurrence starts from the beam entry for P(l=abc @ t=2).
+  //
+  // For this case, the length of the new sequence equals t+1 (t
+  // starts at 0).  This special case can be calculated as:
+  //   P(l=abc? @ t=3) = P(a @ 0)*P(b @ 1)*P(c @ 2)*P(? @ 3)
+  // but we calculate it recursively for speed purposes.
+  typedef ctc_beam_search::BeamEntry<CTCBeamState> BeamEntry;
+  typedef ctc_beam_search::BeamRoot<CTCBeamState> BeamRoot;
+  typedef ctc_beam_search::BeamProbability BeamProbability;
+
+ public:
+  typedef BaseBeamScorer<CTCBeamState> DefaultBeamScorer;
+
+  // The beam search decoder is constructed specifying the beam_width (number of
+  // candidates to keep at each decoding timestep) and a beam scorer (used for
+  // custom scoring, for example enabling the use of a language model).
+  // The ownership of the scorer remains with the caller. The default
+  // implementation, CTCBeamSearchDecoder<>::DefaultBeamScorer, generates the
+  // standard beam search.
+  CTCBeamSearchDecoder(int num_classes, int beam_width,
+                       BaseBeamScorer<CTCBeamState>* scorer, int batch_size = 1,
+                       bool merge_repeated = false)
+      : CTCDecoder(num_classes, batch_size, merge_repeated),
+        beam_width_(beam_width),
+        leaves_(beam_width),
+        beam_scorer_(scorer) {
+    Reset();
+  }
+
+  ~CTCBeamSearchDecoder() override {}
+
+  // Run the hibernating beam search algorithm on the given input.
+  bool Decode(const CTCDecoder::SequenceLength& seq_len,
+              const std::vector<CTCDecoder::Input>& input,
+              std::vector<CTCDecoder::Output>* output,
+              CTCDecoder::ScoreOutput* scores) override;
+
+  // Calculate the next step of the beam search and update the internal state.
+  template <typename Vector>
+  void Step(const Vector& raw_input);
+
+  template <typename Vector>
+  float GetTopK(const int K, const Vector& input,
+                std::vector<float>* top_k_logits,
+                std::vector<int>* top_k_indices);
+
+  // Retrieve the beam scorer instance used during decoding.
+  BaseBeamScorer<CTCBeamState>* GetBeamScorer() const { return beam_scorer_; }
+
+  // Set label selection parameters for faster decoding.
+  // See comments for label_selection_size_ and label_selection_margin_.
+  void SetLabelSelectionParameters(int label_selection_size,
+                                   float label_selection_margin) {
+    label_selection_size_ = label_selection_size;
+    label_selection_margin_ = label_selection_margin;
+  }
+
+  // Reset the beam search
+  void Reset();
+
+  // Extract the top n paths at current time step
+  bool TopPaths(int n, std::vector<std::vector<int>>* paths,
+                std::vector<float>* log_probs, bool merge_repeated) const;
+
+ private:
+  int beam_width_;
+
+  // Label selection is designed to avoid possibly very expensive scorer calls,
+  // by pruning the hypotheses based on the input alone.
+  // Label selection size controls how many items in each beam are passed
+  // through to the beam scorer. Only items with top N input scores are
+  // considered.
+  // Label selection margin controls the difference between minimal input score
+  // (versus the best scoring label) for an item to be passed to the beam
+  // scorer. This margin is expressed in terms of log-probability.
+  // Default is to do no label selection.
+  // For more detail: https://research.google.com/pubs/pub44823.html
+  int label_selection_size_ = 0;       // zero means unlimited
+  float label_selection_margin_ = -1;  // -1 means unlimited.
+
+  gtl::TopN<BeamEntry*, CTCBeamComparer> leaves_;
+  std::unique_ptr<BeamRoot> beam_root_;
+  BaseBeamScorer<CTCBeamState>* beam_scorer_;
+
+  CTCBeamSearchDecoder(const CTCBeamSearchDecoder&) = delete;
+  void operator=(const CTCBeamSearchDecoder&) = delete;
+};
+
+template <typename CTCBeamState, typename CTCBeamComparer>
+bool CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::Decode(
+    const CTCDecoder::SequenceLength& seq_len,
+    const std::vector<CTCDecoder::Input>& input,
+    std::vector<CTCDecoder::Output>* output, ScoreOutput* scores) {
+  // Storage for top paths.
+  std::vector<std::vector<int>> beams;
+  std::vector<float> beam_log_probabilities;
+  int top_n = output->size();
+  if (std::any_of(output->begin(), output->end(),
+                  [this](const CTCDecoder::Output& output) -> bool {
+                    return output.size() < this->batch_size_;
+                  })) {
+    return false;
+  }
+  if (scores->rows() < batch_size_ || scores->cols() < top_n) {
+    return false;
+  }
+
+  for (int b = 0; b < batch_size_; ++b) {
+    int seq_len_b = seq_len[b];
+    Reset();
+
+    for (int t = 0; t < seq_len_b; ++t) {
+      // Pass log-probabilities for this example + time.
+      Step(input[t].row(b));
+    }  // for (int t...
+
+    // O(n * log(n))
+    std::unique_ptr<std::vector<BeamEntry*>> branches(leaves_.Extract());
+    leaves_.Reset();
+    for (int i = 0; i < branches->size(); ++i) {
+      BeamEntry* entry = (*branches)[i];
+      beam_scorer_->ExpandStateEnd(&entry->state);
+      entry->newp.total +=
+          beam_scorer_->GetStateEndExpansionScore(entry->state);
+      leaves_.push(entry);
+    }
+
+    bool status =
+        TopPaths(top_n, &beams, &beam_log_probabilities, merge_repeated_);
+    if (!status) {
+      return status;
+    }
+
+    TFLITE_DCHECK_EQ(top_n, beam_log_probabilities.size());
+    TFLITE_DCHECK_EQ(beams.size(), beam_log_probabilities.size());
+
+    for (int i = 0; i < top_n; ++i) {
+      // Copy output to the correct beam + batch
+      (*output)[i][b].swap(beams[i]);
+      (*scores)(b, i) = -beam_log_probabilities[i];
+    }
+  }  // for (int b...
+  return true;
+}
+
+template <typename CTCBeamState, typename CTCBeamComparer>
+template <typename Vector>
+float CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::GetTopK(
+    const int K, const Vector& input, std::vector<float>* top_k_logits,
+    std::vector<int>* top_k_indices) {
+  // Find Top K choices, complexity nk in worst case. The array input is read
+  // just once.
+  TFLITE_DCHECK_EQ(num_classes_, input.size());
+  top_k_logits->clear();
+  top_k_indices->clear();
+  top_k_logits->resize(K, -INFINITY);
+  top_k_indices->resize(K, -1);
+  for (int j = 0; j < num_classes_ - 1; ++j) {
+    const float logit = input(j);
+    if (logit > (*top_k_logits)[K - 1]) {
+      int k = K - 1;
+      while (k > 0 && logit > (*top_k_logits)[k - 1]) {
+        (*top_k_logits)[k] = (*top_k_logits)[k - 1];
+        (*top_k_indices)[k] = (*top_k_indices)[k - 1];
+        k--;
+      }
+      (*top_k_logits)[k] = logit;
+      (*top_k_indices)[k] = j;
+    }
+  }
+  // Return max value which is in 0th index or blank character logit
+  return std::max((*top_k_logits)[0], input(num_classes_ - 1));
+}
+
+template <typename CTCBeamState, typename CTCBeamComparer>
+template <typename Vector>
+void CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::Step(
+    const Vector& raw_input) {
+  std::vector<float> top_k_logits;
+  std::vector<int> top_k_indices;
+  const bool top_k =
+      (label_selection_size_ > 0 && label_selection_size_ < raw_input.size());
+  // Number of character classes to consider in each step.
+  const int max_classes = top_k ? label_selection_size_ : (num_classes_ - 1);
+  // Get max coefficient and remove it from raw_input later.
+  float max_coeff;
+  if (top_k) {
+    max_coeff = GetTopK(label_selection_size_, raw_input, &top_k_logits,
+                        &top_k_indices);
+  } else {
+    max_coeff = raw_input.maxCoeff();
+  }
+
+  // Get normalization term of softmax: log(sum(exp(logit[j]-max_coeff))).
+  float logsumexp = 0.0;
+  for (int j = 0; j < raw_input.size(); ++j) {
+    logsumexp += Eigen::numext::exp(raw_input(j) - max_coeff);
+  }
+  logsumexp = Eigen::numext::log(logsumexp);
+  // Final normalization offset to get correct log probabilities.
+  float norm_offset = max_coeff + logsumexp;
+
+  const float label_selection_input_min =
+      (label_selection_margin_ >= 0) ? (max_coeff - label_selection_margin_)
+                                     : -std::numeric_limits<float>::infinity();
+
+  // Extract the beams sorted in decreasing new probability
+  TFLITE_DCHECK_EQ(num_classes_, raw_input.size());
+
+  std::unique_ptr<std::vector<BeamEntry*>> branches(leaves_.Extract());
+  leaves_.Reset();
+
+  for (BeamEntry* b : *branches) {
+    // P(.. @ t) becomes the new P(.. @ t-1)
+    b->oldp = b->newp;
+  }
+
+  for (BeamEntry* b : *branches) {
+    if (b->parent != nullptr) {  // if not the root
+      if (b->parent->Active()) {
+        // If last two sequence characters are identical:
+        //   Plabel(l=acc @ t=6) = (Plabel(l=acc @ t=5)
+        //                          + Pblank(l=ac @ t=5))
+        // else:
+        //   Plabel(l=abc @ t=6) = (Plabel(l=abc @ t=5)
+        //                          + P(l=ab @ t=5))
+        float previous = (b->label == b->parent->label) ? b->parent->oldp.blank
+                                                        : b->parent->oldp.total;
+        b->newp.label =
+            LogSumExp(b->newp.label,
+                      beam_scorer_->GetStateExpansionScore(b->state, previous));
+      }
+      // Plabel(l=abc @ t=6) *= P(c @ 6)
+      b->newp.label += raw_input(b->label) - norm_offset;
+    }
+    // Pblank(l=abc @ t=6) = P(l=abc @ t=5) * P(- @ 6)
+    b->newp.blank = b->oldp.total + raw_input(blank_index_) - norm_offset;
+    // P(l=abc @ t=6) = Plabel(l=abc @ t=6) + Pblank(l=abc @ t=6)
+    b->newp.total = LogSumExp(b->newp.blank, b->newp.label);
+
+    // Push the entry back to the top paths list.
+    // Note, this will always fill leaves back up in sorted order.
+    leaves_.push(b);
+  }
+
+  // we need to resort branches in descending oldp order.
+
+  // branches is in descending oldp order because it was
+  // originally in descending newp order and we copied newp to oldp.
+
+  // Grow new leaves
+  for (BeamEntry* b : *branches) {
+    // A new leaf (represented by its BeamProbability) is a candidate
+    // iff its total probability is nonzero and either the beam list
+    // isn't full, or the lowest probability entry in the beam has a
+    // lower probability than the leaf.
+    auto is_candidate = [this](const BeamProbability& prob) {
+      return (prob.total > kLogZero &&
+              (leaves_.size() < beam_width_ ||
+               prob.total > leaves_.peek_bottom()->newp.total));
+    };
+
+    if (!is_candidate(b->oldp)) {
+      continue;
+    }
+
+    for (int ind = 0; ind < max_classes; ind++) {
+      const int label = top_k ? top_k_indices[ind] : ind;
+      const float logit = top_k ? top_k_logits[ind] : raw_input(ind);
+      // Perform label selection: if input for this label looks very
+      // unpromising, never evaluate it with a scorer.
+      // We may compare logits instead of log probabilities,
+      // since the difference is the same in both cases.
+      if (logit < label_selection_input_min) {
+        continue;
+      }
+      BeamEntry& c = b->GetChild(label);
+      if (!c.Active()) {
+        //   Pblank(l=abcd @ t=6) = 0
+        c.newp.blank = kLogZero;
+        // If new child label is identical to beam label:
+        //   Plabel(l=abcc @ t=6) = Pblank(l=abc @ t=5) * P(c @ 6)
+        // Otherwise:
+        //   Plabel(l=abcd @ t=6) = P(l=abc @ t=5) * P(d @ 6)
+        beam_scorer_->ExpandState(b->state, b->label, &c.state, c.label);
+        float previous = (c.label == b->label) ? b->oldp.blank : b->oldp.total;
+        c.newp.label = logit - norm_offset +
+                       beam_scorer_->GetStateExpansionScore(c.state, previous);
+        // P(l=abcd @ t=6) = Plabel(l=abcd @ t=6)
+        c.newp.total = c.newp.label;
+
+        if (is_candidate(c.newp)) {
+          // Before adding the new node to the beam, check if the beam
+          // is already at maximum width.
+          if (leaves_.size() == beam_width_) {
+            // Bottom is no longer in the beam search.  Reset
+            // its probability; signal it's no longer in the beam search.
+            BeamEntry* bottom = leaves_.peek_bottom();
+            bottom->newp.Reset();
+          }
+          leaves_.push(&c);
+        } else {
+          // Deactivate child.
+          c.oldp.Reset();
+          c.newp.Reset();
+        }
+      }
+    }
+  }  // for (BeamEntry* b...
+}
+
+template <typename CTCBeamState, typename CTCBeamComparer>
+void CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::Reset() {
+  leaves_.Reset();
+
+  // This beam root, and all of its children, will be in memory until
+  // the next reset.
+  beam_root_.reset(new BeamRoot(nullptr, -1));
+  beam_root_->RootEntry()->newp.total = 0.0;  // ln(1)
+  beam_root_->RootEntry()->newp.blank = 0.0;  // ln(1)
+
+  // Add the root as the initial leaf.
+  leaves_.push(beam_root_->RootEntry());
+
+  // Call initialize state on the root object.
+  beam_scorer_->InitializeState(&beam_root_->RootEntry()->state);
+}
+
+template <typename CTCBeamState, typename CTCBeamComparer>
+bool CTCBeamSearchDecoder<CTCBeamState, CTCBeamComparer>::TopPaths(
+    int n, std::vector<std::vector<int>>* paths, std::vector<float>* log_probs,
+    bool merge_repeated) const {
+  TFLITE_DCHECK(paths);
+  TFLITE_DCHECK(log_probs);
+  paths->clear();
+  log_probs->clear();
+  if (n > beam_width_) {
+    return false;
+  }
+  if (n > leaves_.size()) {
+    return false;
+  }
+
+  gtl::TopN<BeamEntry*, CTCBeamComparer> top_branches(n);
+
+  // O(beam_width_ * log(n)), space complexity is O(n)
+  for (auto it = leaves_.unsorted_begin(); it != leaves_.unsorted_end(); ++it) {
+    top_branches.push(*it);
+  }
+  // O(n * log(n))
+  std::unique_ptr<std::vector<BeamEntry*>> branches(top_branches.Extract());
+
+  for (int i = 0; i < n; ++i) {
+    BeamEntry* e((*branches)[i]);
+    paths->push_back(e->LabelSeq(merge_repeated));
+    log_probs->push_back(e->newp.total);
+  }
+  return true;
+}
+
+}  // namespace ctc
+}  // namespace custom
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_CTC_BEAM_SEARCH_H_

tensorflow/lite/kernels/ctc/ctc_beam_search_decoder.cc

@@ -0,0 +1,272 @@
+/* Copyright 2018 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 <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <initializer_list>
+#include <vector>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/ctc/ctc_beam_search.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+namespace ctc_beam_search_decoder {
+
+constexpr int kInputsTensor = 0;
+constexpr int kSequenceLengthTensor = 1;
+
+typedef struct {
+  int beam_width;
+  int top_paths;
+  bool merge_repeated;
+} CTCBeamSearchDecoderParams;
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  TFLITE_CHECK(buffer != nullptr);
+  const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
+  const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap();
+
+  CTCBeamSearchDecoderParams* option = new CTCBeamSearchDecoderParams;
+  option->beam_width = m["beam_width"].AsInt32();
+  option->top_paths = m["top_paths"].AsInt32();
+  option->merge_repeated = m["merge_repeated"].AsBool();
+
+  return option;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<CTCBeamSearchDecoderParams*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  const CTCBeamSearchDecoderParams* option =
+      reinterpret_cast<CTCBeamSearchDecoderParams*>(node->user_data);
+  const int top_paths = option->top_paths;
+  TF_LITE_ENSURE(context, option->beam_width >= top_paths);
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  // The outputs should be top_paths * 3 + 1.
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 3 * top_paths + 1);
+
+  const TfLiteTensor* inputs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputsTensor, &inputs));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(inputs), 3);
+  // TensorFlow only supports float.
+  TF_LITE_ENSURE_EQ(context, inputs->type, kTfLiteFloat32);
+  const int batch_size = SizeOfDimension(inputs, 1);
+
+  const TfLiteTensor* sequence_length;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
+                                          &sequence_length));
+  TF_LITE_ENSURE_EQ(context, NumDimensions(sequence_length), 1);
+  TF_LITE_ENSURE_EQ(context, NumElements(sequence_length), batch_size);
+  // TensorFlow only supports int32.
+  TF_LITE_ENSURE_EQ(context, sequence_length->type, kTfLiteInt32);
+
+  // Resize decoded outputs.
+  // Do not resize indices & values cause we don't know the values yet.
+  for (int i = 0; i < top_paths; ++i) {
+    TfLiteTensor* indices;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i, &indices));
+    SetTensorToDynamic(indices);
+    TfLiteTensor* values;
+    TF_LITE_ENSURE_OK(context,
+                      GetOutputSafe(context, node, i + top_paths, &values));
+    SetTensorToDynamic(values);
+    TfLiteTensor* output_shape;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i + 2 * top_paths,
+                                             &output_shape));
+    SetTensorToDynamic(output_shape);
+  }
+
+  // Resize log probability outputs.
+  TfLiteTensor* log_probability_output;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, top_paths * 3,
+                                           &log_probability_output));
+  TfLiteIntArray* log_probability_output_shape_array = TfLiteIntArrayCreate(2);
+  log_probability_output_shape_array->data[0] = batch_size;
+  log_probability_output_shape_array->data[1] = top_paths;
+  return context->ResizeTensor(context, log_probability_output,
+                               log_probability_output_shape_array);
+}
+
+TfLiteStatus Resize(TfLiteContext* context,
+                    std::initializer_list<int32_t> output_shape,
+                    TfLiteTensor* output) {
+  const int dimensions = output_shape.size();
+  TfLiteIntArray* output_shape_array = TfLiteIntArrayCreate(dimensions);
+  int i = 0;
+  for (const int v : output_shape) {
+    output_shape_array->data[i++] = v;
+  }
+  return context->ResizeTensor(context, output, output_shape_array);
+}
+
+TfLiteStatus StoreAllDecodedSequences(
+    TfLiteContext* context,
+    const std::vector<std::vector<std::vector<int>>>& sequences,
+    TfLiteNode* node, int top_paths) {
+  const int32_t batch_size = sequences.size();
+  std::vector<int32_t> num_entries(top_paths, 0);
+
+  // Calculate num_entries per path
+  for (const auto& batch_s : sequences) {
+    TF_LITE_ENSURE_EQ(context, batch_s.size(), top_paths);
+    for (int p = 0; p < top_paths; ++p) {
+      num_entries[p] += batch_s[p].size();
+    }
+  }
+
+  for (int p = 0; p < top_paths; ++p) {
+    const int32_t p_num = num_entries[p];
+
+    // Resize the decoded outputs.
+    TfLiteTensor* indices;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p, &indices));
+    TF_LITE_ENSURE_OK(context, Resize(context, {p_num, 2}, indices));
+
+    TfLiteTensor* values;
+    TF_LITE_ENSURE_OK(context,
+                      GetOutputSafe(context, node, p + top_paths, &values));
+    TF_LITE_ENSURE_OK(context, Resize(context, {p_num}, values));
+
+    TfLiteTensor* decoded_shape;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, p + 2 * top_paths,
+                                             &decoded_shape));
+    TF_LITE_ENSURE_OK(context, Resize(context, {2}, decoded_shape));
+
+    int32_t max_decoded = 0;
+    int32_t offset = 0;
+
+    int32_t* indices_data = GetTensorData<int32_t>(indices);
+    int32_t* values_data = GetTensorData<int32_t>(values);
+    int32_t* decoded_shape_data = GetTensorData<int32_t>(decoded_shape);
+    for (int b = 0; b < batch_size; ++b) {
+      auto& p_batch = sequences[b][p];
+      int32_t num_decoded = p_batch.size();
+      max_decoded = std::max(max_decoded, num_decoded);
+
+      std::copy_n(p_batch.begin(), num_decoded, values_data + offset);
+      for (int32_t t = 0; t < num_decoded; ++t, ++offset) {
+        indices_data[offset * 2] = b;
+        indices_data[offset * 2 + 1] = t;
+      }
+    }
+
+    decoded_shape_data[0] = batch_size;
+    decoded_shape_data[1] = max_decoded;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* inputs;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kInputsTensor, &inputs));
+  const TfLiteTensor* sequence_length;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kSequenceLengthTensor,
+                                          &sequence_length));
+  const CTCBeamSearchDecoderParams* option =
+      reinterpret_cast<CTCBeamSearchDecoderParams*>(node->user_data);
+
+  const int max_time = SizeOfDimension(inputs, 0);
+  const int batch_size = SizeOfDimension(inputs, 1);
+  const int num_classes = SizeOfDimension(inputs, 2);
+
+  const int beam_width = option->beam_width;
+  const int top_paths = option->top_paths;
+  const bool merge_repeated = option->merge_repeated;
+
+  // Validate sequence length is less or equal than max time.
+  for (int i = 0; i < batch_size; ++i) {
+    TF_LITE_ENSURE(context,
+                   max_time >= GetTensorData<int32_t>(sequence_length)[i]);
+  }
+
+  // The following logic is implemented like
+  // tensorflow/core/kernels/ctc_decoder_ops.cc
+  std::vector<optimized_ops::TTypes<float>::UnalignedConstMatrix> input_list_t;
+
+  input_list_t.reserve(max_time);
+  for (std::size_t t = 0; t < max_time; ++t) {
+    input_list_t.emplace_back(
+        GetTensorData<float>(inputs) + t * batch_size * num_classes, batch_size,
+        num_classes);
+  }
+
+  ::tflite::custom::ctc::CTCBeamSearchDecoder<>::DefaultBeamScorer beam_scorer;
+  ::tflite::custom::ctc::CTCBeamSearchDecoder<> beam_search(
+      num_classes, beam_width, &beam_scorer, 1 /* batch_size */,
+      merge_repeated);
+
+  // Allocate temporary memory for holding chip operation data.
+  float* input_chip_t_data =
+      static_cast<float*>(malloc(num_classes * sizeof(float)));
+  Eigen::array<Eigen::DenseIndex, 1> dims;
+  dims[0] = num_classes;
+  optimized_ops::TTypes<float>::Flat input_chip_t(input_chip_t_data, dims);
+
+  std::vector<std::vector<std::vector<int>>> best_paths(batch_size);
+  std::vector<float> log_probs;
+
+  TfLiteTensor* log_probabilities;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node, 3 * top_paths, &log_probabilities));
+  float* log_probabilities_output = GetTensorData<float>(log_probabilities);
+
+  // Assumption: the blank index is num_classes - 1
+  for (int b = 0; b < batch_size; ++b) {
+    auto& best_paths_b = best_paths[b];
+    best_paths_b.resize(top_paths);
+    for (int t = 0; t < GetTensorData<int32_t>(sequence_length)[b]; ++t) {
+      input_chip_t = input_list_t[t].chip(b, 0);
+      auto input_bi =
+          Eigen::Map<const Eigen::ArrayXf>(input_chip_t.data(), num_classes);
+      beam_search.Step(input_bi);
+    }
+    TF_LITE_ENSURE(context, beam_search.TopPaths(top_paths, &best_paths_b,
+                                                 &log_probs, merge_repeated));
+    beam_search.Reset();
+
+    // Fill in log_probabilities output.
+    for (int bp = 0; bp < top_paths; ++bp) {
+      log_probabilities_output[b * top_paths + bp] = log_probs[bp];
+    }
+  }
+
+  free(input_chip_t_data);
+  return StoreAllDecodedSequences(context, best_paths, node, top_paths);
+}
+
+}  // namespace ctc_beam_search_decoder
+
+TfLiteRegistration* Register_CTC_BEAM_SEARCH_DECODER() {
+  static TfLiteRegistration r = {
+      ctc_beam_search_decoder::Init, ctc_beam_search_decoder::Free,
+      ctc_beam_search_decoder::Prepare, ctc_beam_search_decoder::Eval};
+  return &r;
+}
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/ctc/ctc_beam_search_decoder_test.cc

@@ -0,0 +1,233 @@
+/* Copyright 2018 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 <functional>
+#include <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+TfLiteRegistration* Register_CTC_BEAM_SEARCH_DECODER();
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+class CTCBeamSearchDecoderOpModel : public SingleOpModel {
+ public:
+  CTCBeamSearchDecoderOpModel(std::initializer_list<int> input_shape,
+                              std::initializer_list<int> sequence_length_shape,
+                              int beam_width, int top_paths,
+                              bool merge_repeated) {
+    inputs_ = AddInput(TensorType_FLOAT32);
+    sequence_length_ = AddInput(TensorType_INT32);
+
+    for (int i = 0; i < top_paths * 3; ++i) {
+      outputs_.push_back(AddOutput(TensorType_INT32));
+    }
+    outputs_.push_back(AddOutput(TensorType_FLOAT32));
+
+    flexbuffers::Builder fbb;
+    fbb.Map([&]() {
+      fbb.Int("beam_width", beam_width);
+      fbb.Int("top_paths", top_paths);
+      fbb.Bool("merge_repeated", merge_repeated);
+    });
+    fbb.Finish();
+    SetCustomOp("CTCBeamSearchDecoder", fbb.GetBuffer(),
+                Register_CTC_BEAM_SEARCH_DECODER);
+    BuildInterpreter({input_shape, sequence_length_shape});
+  }
+
+  int inputs() { return inputs_; }
+
+  int sequence_length() { return sequence_length_; }
+
+  std::vector<std::vector<int>> GetDecodedOutpus() {
+    std::vector<std::vector<int>> outputs;
+    for (int i = 0; i < outputs_.size() - 1; ++i) {
+      outputs.push_back(ExtractVector<int>(outputs_[i]));
+    }
+    return outputs;
+  }
+
+  std::vector<float> GetLogProbabilitiesOutput() {
+    return ExtractVector<float>(outputs_[outputs_.size() - 1]);
+  }
+
+  std::vector<std::vector<int>> GetOutputShapes() {
+    std::vector<std::vector<int>> output_shapes;
+    for (const int output : outputs_) {
+      output_shapes.push_back(GetTensorShape(output));
+    }
+    return output_shapes;
+  }
+
+ private:
+  int inputs_;
+  int sequence_length_;
+  std::vector<int> outputs_;
+};
+
+TEST(CTCBeamSearchTest, SimpleTest) {
+  CTCBeamSearchDecoderOpModel m({2, 1, 2}, {1}, 1, 1, true);
+  m.PopulateTensor<float>(m.inputs(),
+                          {-0.50922557, -1.35512652, -2.55445064, -1.58419356});
+  m.PopulateTensor<int>(m.sequence_length(), {2});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  // Make sure the output shapes are right.
+  const std::vector<std::vector<int>>& output_shapes = m.GetOutputShapes();
+  EXPECT_EQ(output_shapes.size(), 4);
+  EXPECT_THAT(output_shapes[0], ElementsAre(1, 2));
+  EXPECT_THAT(output_shapes[1], ElementsAre(1));
+  EXPECT_THAT(output_shapes[2], ElementsAre(2));
+  EXPECT_THAT(output_shapes[3], ElementsAre(1, 1));
+
+  // Check decoded outputs.
+  const std::vector<std::vector<int>>& decoded_outputs = m.GetDecodedOutpus();
+  EXPECT_EQ(decoded_outputs.size(), 3);
+  EXPECT_THAT(decoded_outputs[0], ElementsAre(0, 0));
+  EXPECT_THAT(decoded_outputs[1], ElementsAre(0));
+  EXPECT_THAT(decoded_outputs[2], ElementsAre(1, 1));
+  // Check log probabilities output.
+  EXPECT_THAT(m.GetLogProbabilitiesOutput(),
+              ElementsAreArray(ArrayFloatNear({-0.357094})));
+}
+
+TEST(CTCBeamSearchTest, MultiBatchTest) {
+  CTCBeamSearchDecoderOpModel m({3, 3, 3}, {3}, 1, 1, true);
+  m.PopulateTensor<float>(
+      m.inputs(),
+      {-0.63649208, -0.00487571, -0.04249819, -0.67754697, -1.0341399,
+       -2.14717721, -0.77686821, -3.41973774, -0.05151402, -0.21482619,
+       -0.57411168, -1.45039917, -0.73769373, -2.10941739, -0.44818325,
+       -0.25287673, -2.80057302, -0.54748312, -0.73334867, -0.86537719,
+       -0.2065197,  -0.18725838, -1.42770405, -0.86051965, -1.61642301,
+       -2.07275114, -0.9201845});
+  m.PopulateTensor<int>(m.sequence_length(), {3, 3, 3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  // Make sure the output shapes are right.
+  const std::vector<std::vector<int>>& output_shapes = m.GetOutputShapes();
+  EXPECT_EQ(output_shapes.size(), 4);
+  EXPECT_THAT(output_shapes[0], ElementsAre(4, 2));
+  EXPECT_THAT(output_shapes[1], ElementsAre(4));
+  EXPECT_THAT(output_shapes[2], ElementsAre(2));
+  EXPECT_THAT(output_shapes[3], ElementsAre(3, 1));
+
+  // Check decoded outputs.
+  const std::vector<std::vector<int>>& decoded_outputs = m.GetDecodedOutpus();
+  EXPECT_EQ(decoded_outputs.size(), 3);
+  EXPECT_THAT(decoded_outputs[0], ElementsAre(0, 0, 0, 1, 1, 0, 2, 0));
+  EXPECT_THAT(decoded_outputs[1], ElementsAre(1, 0, 0, 0));
+  EXPECT_THAT(decoded_outputs[2], ElementsAre(3, 2));
+  // Check log probabilities output.
+  EXPECT_THAT(m.GetLogProbabilitiesOutput(),
+              ElementsAreArray(ArrayFloatNear({-1.88343, -1.41188, -1.20958})));
+}
+
+TEST(CTCBeamSearchTest, MultiPathsTest) {
+  CTCBeamSearchDecoderOpModel m({3, 2, 5}, {2}, 3, 2, true);
+  m.PopulateTensor<float>(
+      m.inputs(),
+      {-2.206851,   -0.09542714, -0.2393415,  -3.81866197, -0.27241158,
+       -0.20371124, -0.68236623, -1.1397166,  -0.17422639, -1.85224048,
+       -0.9406037,  -0.32544678, -0.21846784, -0.38377237, -0.33498676,
+       -0.10139782, -0.51886883, -0.21678554, -0.15267063, -1.91164412,
+       -0.31328673, -0.27462716, -0.65975336, -1.53671973, -2.76554225,
+       -0.23920634, -1.2370502,  -4.98751576, -3.12995717, -0.43129368});
+  m.PopulateTensor<int>(m.sequence_length(), {3, 3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  // Make sure the output shapes are right.
+  const std::vector<std::vector<int>>& output_shapes = m.GetOutputShapes();
+  EXPECT_EQ(output_shapes.size(), 7);
+  EXPECT_THAT(output_shapes[0], ElementsAre(4, 2));
+  EXPECT_THAT(output_shapes[1], ElementsAre(3, 2));
+  EXPECT_THAT(output_shapes[2], ElementsAre(4));
+  EXPECT_THAT(output_shapes[3], ElementsAre(3));
+  EXPECT_THAT(output_shapes[4], ElementsAre(2));
+  EXPECT_THAT(output_shapes[5], ElementsAre(2));
+  EXPECT_THAT(output_shapes[6], ElementsAre(2, 2));
+
+  // Check decoded outputs.
+  const std::vector<std::vector<int>>& decoded_outputs = m.GetDecodedOutpus();
+  EXPECT_EQ(decoded_outputs.size(), 6);
+  EXPECT_THAT(decoded_outputs[0], ElementsAre(0, 0, 0, 1, 1, 0, 1, 1));
+  EXPECT_THAT(decoded_outputs[1], ElementsAre(0, 0, 0, 1, 1, 0));
+  EXPECT_THAT(decoded_outputs[2], ElementsAre(1, 2, 3, 0));
+  EXPECT_THAT(decoded_outputs[3], ElementsAre(2, 1, 0));
+  EXPECT_THAT(decoded_outputs[4], ElementsAre(2, 2));
+  EXPECT_THAT(decoded_outputs[5], ElementsAre(2, 2));
+  // Check log probabilities output.
+  EXPECT_THAT(m.GetLogProbabilitiesOutput(),
+              ElementsAreArray(
+                  ArrayFloatNear({-2.65148, -2.65864, -2.17914, -2.61357})));
+}
+
+TEST(CTCBeamSearchTest, NonEqualSequencesTest) {
+  CTCBeamSearchDecoderOpModel m({3, 3, 4}, {3}, 3, 1, true);
+  m.PopulateTensor<float>(
+      m.inputs(),
+      {-1.26658163, -0.25760023, -0.03917975, -0.63772235, -0.03794756,
+       -0.45063099, -0.27706473, -0.01569179, -0.59940385, -0.35700127,
+       -0.48920721, -1.42635476, -1.3462478,  -0.02565498, -0.30179568,
+       -0.6491698,  -0.55017719, -2.92291466, -0.92522973, -0.47592022,
+       -0.07099135, -0.31575624, -0.86345281, -0.36017021, -0.79208612,
+       -1.75306124, -0.65089224, -0.00912786, -0.42915003, -1.72606203,
+       -1.66337589, -0.70800793, -2.52272352, -0.67329562, -2.49145522,
+       -0.49786342});
+  m.PopulateTensor<int>(m.sequence_length(), {1, 2, 3});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  // Make sure the output shapes are right.
+  const std::vector<std::vector<int>>& output_shapes = m.GetOutputShapes();
+  EXPECT_EQ(output_shapes.size(), 4);
+  EXPECT_THAT(output_shapes[0], ElementsAre(3, 2));
+  EXPECT_THAT(output_shapes[1], ElementsAre(3));
+  EXPECT_THAT(output_shapes[2], ElementsAre(2));
+  EXPECT_THAT(output_shapes[3], ElementsAre(3, 1));
+
+  // Check decoded outputs.
+  const std::vector<std::vector<int>>& decoded_outputs = m.GetDecodedOutpus();
+  EXPECT_EQ(decoded_outputs.size(), 3);
+  EXPECT_THAT(decoded_outputs[0], ElementsAre(0, 0, 1, 0, 2, 0));
+  EXPECT_THAT(decoded_outputs[1], ElementsAre(2, 0, 1));
+  EXPECT_THAT(decoded_outputs[2], ElementsAre(3, 1));
+  // Check log probabilities output.
+  EXPECT_THAT(m.GetLogProbabilitiesOutput(),
+              ElementsAreArray(ArrayFloatNear({-0.97322, -1.16334, -2.15553})));
+}
+
+}  // namespace
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+

tensorflow/lite/kernels/ctc/ctc_decoder.h

@@ -0,0 +1,114 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Copied from tensorflow/core/util/ctc/ctc_decoder.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_CTC_DECODER_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_CTC_DECODER_H_
+
+#include <memory>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+
+namespace tflite {
+namespace custom {
+namespace ctc {
+
+// The CTCDecoder is an abstract interface to be implemented when providing a
+// decoding method on the timestep output of a RNN trained with CTC loss.
+//
+// The two types of decoding available are:
+//   - greedy path, through the CTCGreedyDecoder
+//   - beam search, through the CTCBeamSearchDecoder
+class CTCDecoder {
+ public:
+  typedef Eigen::Map<const Eigen::ArrayXi> SequenceLength;
+  typedef Eigen::Map<const Eigen::MatrixXf> Input;
+  typedef std::vector<std::vector<int>> Output;
+  typedef Eigen::Map<Eigen::MatrixXf> ScoreOutput;
+
+  CTCDecoder(int num_classes, int batch_size, bool merge_repeated)
+      : num_classes_(num_classes),
+        blank_index_(num_classes - 1),
+        batch_size_(batch_size),
+        merge_repeated_(merge_repeated) {}
+
+  virtual ~CTCDecoder() {}
+
+  // Dimensionality of the input/output is expected to be:
+  //  - seq_len[b] - b = 0 to batch_size_
+  //  - input[t].rows(b) - t = 0 to timesteps; b = 0 t batch_size_
+  //  - output.size() specifies the number of beams to be returned.
+  //  - scores(b, i) - b = 0 to batch_size; i = 0 to output.size()
+  virtual bool Decode(const SequenceLength& seq_len,
+                      const std::vector<Input>& input,
+                      std::vector<Output>* output, ScoreOutput* scores) = 0;
+
+  int batch_size() { return batch_size_; }
+  int num_classes() { return num_classes_; }
+
+ protected:
+  int num_classes_;
+  int blank_index_;
+  int batch_size_;
+  bool merge_repeated_;
+};
+
+// CTCGreedyDecoder is an implementation of the simple best path decoding
+// algorithm, selecting at each timestep the most likely class at each timestep.
+class CTCGreedyDecoder : public CTCDecoder {
+ public:
+  CTCGreedyDecoder(int num_classes, int batch_size, bool merge_repeated)
+      : CTCDecoder(num_classes, batch_size, merge_repeated) {}
+
+  bool Decode(const CTCDecoder::SequenceLength& seq_len,
+              const std::vector<CTCDecoder::Input>& input,
+              std::vector<CTCDecoder::Output>* output,
+              CTCDecoder::ScoreOutput* scores) override {
+    if (output->empty() || (*output)[0].size() < batch_size_) {
+      return false;
+    }
+    if (scores->rows() < batch_size_ || scores->cols() == 0) {
+      return false;
+    }
+    // For each batch entry, identify the transitions
+    for (int b = 0; b < batch_size_; ++b) {
+      int seq_len_b = seq_len[b];
+      // Only writing to beam 0
+      std::vector<int>& output_b = (*output)[0][b];
+
+      int prev_class_ix = -1;
+      (*scores)(b, 0) = 0;
+      for (int t = 0; t < seq_len_b; ++t) {
+        auto row = input[t].row(b);
+        int max_class_ix;
+        (*scores)(b, 0) += -row.maxCoeff(&max_class_ix);
+        if (max_class_ix != blank_index_ &&
+            !(merge_repeated_ && max_class_ix == prev_class_ix)) {
+          output_b.push_back(max_class_ix);
+        }
+        prev_class_ix = max_class_ix;
+      }
+    }
+    return true;
+  }
+};
+
+}  // namespace ctc
+}  // namespace custom
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_CTC_DECODER_H_

tensorflow/lite/kernels/ctc/ctc_loss_util.h

@@ -0,0 +1,50 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// Copied from tensorflow/core/util/ctc/ctc_loss_util.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_CTC_LOSS_UTIL_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_CTC_LOSS_UTIL_H_
+
+#include <cmath>
+#include <limits>
+
+namespace tflite {
+namespace custom {
+namespace ctc {
+
+const float kLogZero = -std::numeric_limits<float>::infinity();
+
+// Add logarithmic probabilities using:
+// ln(a + b) = ln(a) + ln(1 + exp(ln(b) - ln(a)))
+// The two inputs are assumed to be log probabilities.
+// (GravesTh) Eq. 7.18
+inline float LogSumExp(float log_prob_1, float log_prob_2) {
+  // Always have 'b' be the smaller number to avoid the exponential from
+  // blowing up.
+  if (log_prob_1 == kLogZero && log_prob_2 == kLogZero) {
+    return kLogZero;
+  } else {
+    return (log_prob_1 > log_prob_2)
+               ? log_prob_1 + log1pf(expf(log_prob_2 - log_prob_1))
+               : log_prob_2 + log1pf(expf(log_prob_1 - log_prob_2));
+  }
+}
+
+}  // namespace ctc
+}  // namespace custom
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_CTC_LOSS_UTIL_H_

tensorflow/lite/kernels/ctc/top_n.h

@@ -0,0 +1,341 @@
+/* Copyright 2018 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.
+==============================================================================*/
+
+// This simple class finds the top n elements of an incrementally provided set
+// of elements which you push one at a time.  If the number of elements exceeds
+// n, the lowest elements are incrementally dropped.  At the end you get
+// a vector of the top elements sorted in descending order (through Extract() or
+// ExtractNondestructive()), or a vector of the top elements but not sorted
+// (through ExtractUnsorted() or ExtractUnsortedNondestructive()).
+//
+// The value n is specified in the constructor.  If there are p elements pushed
+// altogether:
+//   The total storage requirements are O(min(n, p)) elements
+//   The running time is O(p * log(min(n, p))) comparisons
+// If n is a constant, the total storage required is a constant and the running
+// time is linear in p.
+//
+// NOTE(zhifengc): There is a way to do this in O(min(n, p)) storage and O(p)
+// runtime. The basic idea is to repeatedly fill up a buffer of 2 * n elements,
+// discarding the lowest n elements whenever the buffer is full using a linear-
+// time median algorithm. This may have better performance when the input
+// sequence is partially sorted.
+//
+// NOTE(zhifengc): This class should be redesigned to avoid reallocating a
+// vector for each Extract.
+
+// Copied from tensorflow/core/lib/gtl/top_n.h
+// TODO(b/111524997): Remove this file.
+#ifndef TENSORFLOW_LITE_KERNELS_CTC_TOP_N_H_
+#define TENSORFLOW_LITE_KERNELS_CTC_TOP_N_H_
+
+#include <stddef.h>
+#include <algorithm>
+#include <functional>
+#include <string>
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace gtl {
+
+// Cmp is an stl binary predicate.  Note that Cmp is the "greater" predicate,
+// not the more commonly used "less" predicate.
+//
+// If you use a "less" predicate here, the TopN will pick out the bottom N
+// elements out of the ones passed to it, and it will return them sorted in
+// ascending order.
+//
+// TopN is rule-of-zero copyable and movable if its members are.
+template <class T, class Cmp = std::greater<T> >
+class TopN {
+ public:
+  // The TopN is in one of the three states:
+  //
+  //  o UNORDERED: this is the state an instance is originally in,
+  //    where the elements are completely orderless.
+  //
+  //  o BOTTOM_KNOWN: in this state, we keep the invariant that there
+  //    is at least one element in it, and the lowest element is at
+  //    position 0. The elements in other positions remain
+  //    unsorted. This state is reached if the state was originally
+  //    UNORDERED and a peek_bottom() function call is invoked.
+  //
+  //  o HEAP_SORTED: in this state, the array is kept as a heap and
+  //    there are exactly (limit_+1) elements in the array. This
+  //    state is reached when at least (limit_+1) elements are
+  //    pushed in.
+  //
+  //  The state transition graph is at follows:
+  //
+  //             peek_bottom()                (limit_+1) elements
+  //  UNORDERED --------------> BOTTOM_KNOWN --------------------> HEAP_SORTED
+  //      |                                                           ^
+  //      |                      (limit_+1) elements                  |
+  //      +-----------------------------------------------------------+
+
+  enum State { UNORDERED, BOTTOM_KNOWN, HEAP_SORTED };
+  using UnsortedIterator = typename std::vector<T>::const_iterator;
+
+  // 'limit' is the maximum number of top results to return.
+  explicit TopN(size_t limit) : TopN(limit, Cmp()) {}
+  TopN(size_t limit, const Cmp &cmp) : limit_(limit), cmp_(cmp) {}
+
+  size_t limit() const { return limit_; }
+
+  // Number of elements currently held by this TopN object.  This
+  // will be no greater than 'limit' passed to the constructor.
+  size_t size() const { return std::min(elements_.size(), limit_); }
+
+  bool empty() const { return size() == 0; }
+
+  // If you know how many elements you will push at the time you create the
+  // TopN object, you can call reserve to preallocate the memory that TopN
+  // will need to process all 'n' pushes.  Calling this method is optional.
+  void reserve(size_t n) { elements_.reserve(std::min(n, limit_ + 1)); }
+
+  // Push 'v'.  If the maximum number of elements was exceeded, drop the
+  // lowest element and return it in 'dropped' (if given). If the maximum is not
+  // exceeded, 'dropped' will remain unchanged. 'dropped' may be omitted or
+  // nullptr, in which case it is not filled in.
+  // Requires: T is CopyAssignable, Swappable
+  void push(const T &v) { push(v, nullptr); }
+  void push(const T &v, T *dropped) { PushInternal(v, dropped); }
+
+  // Move overloads of push.
+  // Requires: T is MoveAssignable, Swappable
+  void push(T &&v) {  // NOLINT(build/c++11)
+    push(std::move(v), nullptr);
+  }
+  void push(T &&v, T *dropped) {  // NOLINT(build/c++11)
+    PushInternal(std::move(v), dropped);
+  }
+
+  // Peeks the bottom result without calling Extract()
+  const T &peek_bottom();
+
+  // Extract the elements as a vector sorted in descending order.  The caller
+  // assumes ownership of the vector and must delete it when done.  This is a
+  // destructive operation.  The only method that can be called immediately
+  // after Extract() is Reset().
+  std::vector<T> *Extract();
+
+  // Similar to Extract(), but makes no guarantees the elements are in sorted
+  // order.  As with Extract(), the caller assumes ownership of the vector and
+  // must delete it when done.  This is a destructive operation.  The only
+  // method that can be called immediately after ExtractUnsorted() is Reset().
+  std::vector<T> *ExtractUnsorted();
+
+  // A non-destructive version of Extract(). Copy the elements in a new vector
+  // sorted in descending order and return it.  The caller assumes ownership of
+  // the new vector and must delete it when done.  After calling
+  // ExtractNondestructive(), the caller can continue to push() new elements.
+  std::vector<T> *ExtractNondestructive() const;
+
+  // A non-destructive version of Extract(). Copy the elements to a given
+  // vector sorted in descending order. After calling
+  // ExtractNondestructive(), the caller can continue to push() new elements.
+  // Note:
+  //  1. The given argument must to be allocated.
+  //  2. Any data contained in the vector prior to the call will be deleted
+  //     from it. After the call the vector will contain only the elements
+  //     from the data structure.
+  void ExtractNondestructive(std::vector<T> *output) const;
+
+  // A non-destructive version of ExtractUnsorted(). Copy the elements in a new
+  // vector and return it, with no guarantees the elements are in sorted order.
+  // The caller assumes ownership of the new vector and must delete it when
+  // done.  After calling ExtractUnsortedNondestructive(), the caller can
+  // continue to push() new elements.
+  std::vector<T> *ExtractUnsortedNondestructive() const;
+
+  // A non-destructive version of ExtractUnsorted(). Copy the elements into
+  // a given vector, with no guarantees the elements are in sorted order.
+  // After calling ExtractUnsortedNondestructive(), the caller can continue
+  // to push() new elements.
+  // Note:
+  //  1. The given argument must to be allocated.
+  //  2. Any data contained in the vector prior to the call will be deleted
+  //     from it. After the call the vector will contain only the elements
+  //     from the data structure.
+  void ExtractUnsortedNondestructive(std::vector<T> *output) const;
+
+  // Return an iterator to the beginning (end) of the container,
+  // with no guarantees about the order of iteration. These iterators are
+  // invalidated by mutation of the data structure.
+  UnsortedIterator unsorted_begin() const { return elements_.begin(); }
+  UnsortedIterator unsorted_end() const { return elements_.begin() + size(); }
+
+  // Accessor for comparator template argument.
+  Cmp *comparator() { return &cmp_; }
+
+  // This removes all elements.  If Extract() or ExtractUnsorted() have been
+  // called, this will put it back in an empty but useable state.
+  void Reset();
+
+ private:
+  template <typename U>
+  void PushInternal(U &&v, T *dropped);  // NOLINT(build/c++11)
+
+  // elements_ can be in one of two states:
+  //   elements_.size() <= limit_:  elements_ is an unsorted vector of elements
+  //      pushed so far.
+  //   elements_.size() > limit_:  The last element of elements_ is unused;
+  //      the other elements of elements_ are an stl heap whose size is exactly
+  //      limit_.  In this case elements_.size() is exactly one greater than
+  //      limit_, but don't use "elements_.size() == limit_ + 1" to check for
+  //      that because you'll get a false positive if limit_ == size_t(-1).
+  std::vector<T> elements_;
+  size_t limit_;  // Maximum number of elements to find
+  Cmp cmp_;       // Greater-than comparison function
+  State state_ = UNORDERED;
+};
+
+// ----------------------------------------------------------------------
+// Implementations of non-inline functions
+
+template <class T, class Cmp>
+template <typename U>
+void TopN<T, Cmp>::PushInternal(U &&v, T *dropped) {  // NOLINT(build/c++11)
+  if (limit_ == 0) {
+    if (dropped) *dropped = std::forward<U>(v);  // NOLINT(build/c++11)
+    return;
+  }
+  if (state_ != HEAP_SORTED) {
+    elements_.push_back(std::forward<U>(v));  // NOLINT(build/c++11)
+    if (state_ == UNORDERED || cmp_(elements_.back(), elements_.front())) {
+      // Easy case: we just pushed the new element back
+    } else {
+      // To maintain the BOTTOM_KNOWN state, we need to make sure that
+      // the element at position 0 is always the smallest. So we put
+      // the new element at position 0 and push the original bottom
+      // element in the back.
+      // Warning: this code is subtle.
+      using std::swap;
+      swap(elements_.front(), elements_.back());
+    }
+    if (elements_.size() == limit_ + 1) {
+      // Transition from unsorted vector to a heap.
+      std::make_heap(elements_.begin(), elements_.end(), cmp_);
+      if (dropped) *dropped = std::move(elements_.front());
+      std::pop_heap(elements_.begin(), elements_.end(), cmp_);
+      state_ = HEAP_SORTED;
+    }
+  } else {
+    // Only insert the new element if it is greater than the least element.
+    if (cmp_(v, elements_.front())) {
+      elements_.back() = std::forward<U>(v);  // NOLINT(build/c++11)
+      std::push_heap(elements_.begin(), elements_.end(), cmp_);
+      if (dropped) *dropped = std::move(elements_.front());
+      std::pop_heap(elements_.begin(), elements_.end(), cmp_);
+    } else {
+      if (dropped) *dropped = std::forward<U>(v);  // NOLINT(build/c++11)
+    }
+  }
+}
+
+template <class T, class Cmp>
+const T &TopN<T, Cmp>::peek_bottom() {
+  TFLITE_DCHECK(!empty());
+  if (state_ == UNORDERED) {
+    // We need to do a linear scan to find out the bottom element
+    int min_candidate = 0;
+    for (size_t i = 1; i < elements_.size(); ++i) {
+      if (cmp_(elements_[min_candidate], elements_[i])) {
+        min_candidate = i;
+      }
+    }
+    // By swapping the element at position 0 and the minimal
+    // element, we transition to the BOTTOM_KNOWN state
+    if (min_candidate != 0) {
+      using std::swap;
+      swap(elements_[0], elements_[min_candidate]);
+    }
+    state_ = BOTTOM_KNOWN;
+  }
+  return elements_.front();
+}
+
+template <class T, class Cmp>
+std::vector<T> *TopN<T, Cmp>::Extract() {
+  auto out = new std::vector<T>;
+  out->swap(elements_);
+  if (state_ != HEAP_SORTED) {
+    std::sort(out->begin(), out->end(), cmp_);
+  } else {
+    out->pop_back();
+    std::sort_heap(out->begin(), out->end(), cmp_);
+  }
+  return out;
+}
+
+template <class T, class Cmp>
+std::vector<T> *TopN<T, Cmp>::ExtractUnsorted() {
+  auto out = new std::vector<T>;
+  out->swap(elements_);
+  if (state_ == HEAP_SORTED) {
+    // Remove the limit_+1'th element.
+    out->pop_back();
+  }
+  return out;
+}
+
+template <class T, class Cmp>
+std::vector<T> *TopN<T, Cmp>::ExtractNondestructive() const {
+  auto out = new std::vector<T>;
+  ExtractNondestructive(out);
+  return out;
+}
+
+template <class T, class Cmp>
+void TopN<T, Cmp>::ExtractNondestructive(std::vector<T> *output) const {
+  TFLITE_DCHECK(output);
+  *output = elements_;
+  if (state_ != HEAP_SORTED) {
+    std::sort(output->begin(), output->end(), cmp_);
+  } else {
+    output->pop_back();
+    std::sort_heap(output->begin(), output->end(), cmp_);
+  }
+}
+
+template <class T, class Cmp>
+std::vector<T> *TopN<T, Cmp>::ExtractUnsortedNondestructive() const {
+  auto elements = new std::vector<T>;
+  ExtractUnsortedNondestructive(elements);
+  return elements;
+}
+
+template <class T, class Cmp>
+void TopN<T, Cmp>::ExtractUnsortedNondestructive(std::vector<T> *output) const {
+  TFLITE_DCHECK(output);
+  *output = elements_;
+  if (state_ == HEAP_SORTED) {
+    // Remove the limit_+1'th element.
+    output->pop_back();
+  }
+}
+
+template <class T, class Cmp>
+void TopN<T, Cmp>::Reset() {
+  elements_.clear();
+  state_ = UNORDERED;
+}
+
+}  // namespace gtl
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CTC_TOP_N_H_

tensorflow/lite/kernels/cumsum.cc

@@ -0,0 +1,110 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace cumsum {
+
+static const int kInputTensor = 0;
+static const int kAxisTensor = 1;
+static const int kOutputTensor = 0;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  const TfLiteTensor* axis = GetInput(context, node, kAxisTensor);
+
+  TF_LITE_ENSURE(context, input->type == kTfLiteInt32 ||
+                              input->type == kTfLiteFloat32 ||
+                              input->type == kTfLiteInt64);
+  TF_LITE_ENSURE_EQ(context, axis->type, kTfLiteInt32);
+
+  TF_LITE_ENSURE_EQ(context, NumElements(axis), 1);
+
+  TF_LITE_ENSURE(context, NumDimensions(input) >= 1);
+
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
+  return context->ResizeTensor(context, output, output_shape);
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input = GetInput(context, node, kInputTensor);
+  const TfLiteTensor* axis_tensor = GetInput(context, node, kAxisTensor);
+
+  TfLiteTensor* output = GetOutput(context, node, kOutputTensor);
+
+  auto* params = reinterpret_cast<TfLiteCumsumParams*>(node->builtin_data);
+
+  int axis = *GetTensorData<int>(axis_tensor);
+  if (axis < 0) axis += NumDimensions(input);
+
+  if (axis < 0 || axis >= NumDimensions(input)) {
+    TF_LITE_KERNEL_LOG(context, "Invalid axis: %d", axis);
+    return kTfLiteError;
+  }
+
+  switch (input->type) {
+    case kTfLiteInt32: {
+      optimized_ops::CumSum(GetTensorData<int>(input), GetTensorShape(input),
+                            axis, params->exclusive, params->reverse,
+                            GetTensorData<int>(output));
+      break;
+    }
+    case kTfLiteInt64: {
+      optimized_ops::CumSum(GetTensorData<int64_t>(input),
+                            GetTensorShape(input), axis, params->exclusive,
+                            params->reverse, GetTensorData<int64_t>(output));
+      break;
+    }
+    case kTfLiteFloat32: {
+      optimized_ops::CumSum(GetTensorData<float>(input), GetTensorShape(input),
+                            axis, params->exclusive, params->reverse,
+                            GetTensorData<float>(output));
+      break;
+    }
+    default: {
+      TF_LITE_KERNEL_LOG(
+          context,
+          "Unsupported input type, cumsum only supports int32 & float32.");
+      return kTfLiteError;
+    }
+  }
+
+  return kTfLiteOk;
+}
+
+}  // namespace cumsum
+
+TfLiteRegistration* Register_CUMSUM() {
+  static TfLiteRegistration r = {nullptr, nullptr, cumsum::Prepare,
+                                 cumsum::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/cumsum_test.cc

@@ -0,0 +1,156 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+
+namespace {
+
+template <typename T>
+class CumsumOpModel : public SingleOpModel {
+ public:
+  CumsumOpModel(const TensorData& input, const TensorData& output,
+                bool exclusive, bool reverse) {
+    input_ = AddInput(input);
+    axis_ = AddInput({TensorType_INT32, {1}});
+
+    output_ = AddOutput(output);
+
+    SetBuiltinOp(BuiltinOperator_CUMSUM, BuiltinOptions_CumsumOptions,
+                 CreateCumsumOptions(builder_, exclusive, reverse).Union());
+
+    BuildInterpreter({GetShape(input_), GetShape(axis_)});
+  }
+
+  int input() { return input_; }
+  int axis() { return axis_; }
+
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+
+ private:
+  int input_;
+  int axis_;
+  int output_;
+};
+
+TEST(CumsumOpTest, SimpleIntTest) {
+  CumsumOpModel<int32_t> m({TensorType_INT32, {2, 4}}, {TensorType_INT32, {}},
+                           false, false);
+
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {1});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              testing::ElementsAreArray({1, 3, 6, 10, 5, 11, 18, 26}));
+}
+
+TEST(CumsumOpTest, SimpleInt64Test) {
+  CumsumOpModel<int64_t> m({TensorType_INT64, {2, 4}}, {TensorType_INT64, {}},
+                           false, false);
+
+  m.PopulateTensor<int64_t>(
+      m.input(), {100000000001l, 100000000002l, 100000000003l, 100000000004l,
+                  100000000005l, 100000000006l, 100000000007l, 100000000008l});
+  m.PopulateTensor<int>(m.axis(), {1});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(
+                                 {100000000001l, 200000000003l, 300000000006l,
+                                  400000000010l, 100000000005l, 200000000011l,
+                                  300000000018l, 400000000026l}));
+}
+
+TEST(CumsumOpTest, SimpleIntAxis0Test) {
+  CumsumOpModel<int32_t> m({TensorType_INT32, {2, 4}}, {TensorType_INT32, {}},
+                           false, false);
+
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {0});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              testing::ElementsAreArray({1, 2, 3, 4, 6, 8, 10, 12}));
+}
+
+TEST(CumsumOpTest, Simple1DIntTest) {
+  CumsumOpModel<int32_t> m({TensorType_INT32, {8}}, {TensorType_INT32, {}},
+                           false, false);
+
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {0});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              testing::ElementsAreArray({1, 3, 6, 10, 15, 21, 28, 36}));
+}
+
+TEST(CumsumOpTest, SimpleIntReverseTest) {
+  CumsumOpModel<int32_t> m({TensorType_INT32, {2, 4}}, {TensorType_INT32, {}},
+                           false, true);
+
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {1});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              testing::ElementsAreArray({10, 9, 7, 4, 26, 21, 15, 8}));
+}
+
+TEST(CumsumOpTest, SimpleIntExclusiveTest) {
+  CumsumOpModel<int32_t> m({TensorType_INT32, {2, 4}}, {TensorType_INT32, {}},
+                           true, false);
+
+  m.PopulateTensor<int>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {1});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              testing::ElementsAreArray({0, 1, 3, 6, 0, 5, 11, 18}));
+}
+
+TEST(CumsumOpTest, SimpleFloatTest) {
+  CumsumOpModel<float> m({TensorType_FLOAT32, {2, 4}}, {TensorType_FLOAT32, {}},
+                         false, false);
+
+  m.PopulateTensor<float>(m.input(), {1, 2, 3, 4, 5, 6, 7, 8});
+  m.PopulateTensor<int>(m.axis(), {1});
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(), testing::ElementsAreArray(
+                                 ArrayFloatNear({1, 3, 6, 10, 5, 11, 18, 26})));
+}
+
+}  // namespace
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/custom_ops_register.h

@@ -0,0 +1,44 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_CUSTOM_OPS_REGISTER_H_
+#define TENSORFLOW_LITE_KERNELS_CUSTOM_OPS_REGISTER_H_
+
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace ops {
+namespace custom {
+
+TfLiteRegistration* Register_ATAN2();
+TfLiteRegistration* Register_AVG_POOL_3D();
+TfLiteRegistration* Register_HASHTABLE();
+TfLiteRegistration* Register_HASHTABLE_FIND();
+TfLiteRegistration* Register_HASHTABLE_IMPORT();
+TfLiteRegistration* Register_HASHTABLE_SIZE();
+TfLiteRegistration* Register_IRFFT2D();
+TfLiteRegistration* Register_MAX_POOL_3D();
+TfLiteRegistration* Register_MULTINOMIAL();
+TfLiteRegistration* Register_RANDOM_STANDARD_NORMAL();
+TfLiteRegistration* Register_RANDOM_UNIFORM();
+TfLiteRegistration* Register_RANDOM_UNIFORM_INT();
+TfLiteRegistration* Register_ROLL();
+TfLiteRegistration* Register_SIGN();
+TfLiteRegistration* Register_TABLE();
+
+}  // namespace custom
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_CUSTOM_OPS_REGISTER_H_

tensorflow/lite/kernels/densify.cc

@@ -0,0 +1,122 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/densify.h"
+
+#include <stddef.h>
+
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace densify {
+
+struct OpContext {
+  OpContext(TfLiteContext* context, TfLiteNode* node) {
+    input = GetInput(context, node, 0);
+    output = GetOutput(context, node, 0);
+  }
+  const TfLiteTensor* input;
+  TfLiteTensor* output;
+};
+
+struct OpData {
+  bool dense_weights_initialized;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  auto* op_data = new OpData();
+  op_data->dense_weights_initialized = false;
+  return op_data;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  OpContext op_context(context, node);
+
+  TF_LITE_ENSURE(context, op_context.input != nullptr);
+  TF_LITE_ENSURE(context, op_context.input->type != kTfLiteString);
+  TF_LITE_ENSURE(context, IsConstantTensor(op_context.input));
+  TF_LITE_ENSURE(context, op_context.input->sparsity != nullptr);
+
+  op_context.output->type = op_context.input->type;
+  op_context.output->name = "Densify_output";
+  op_context.output->allocation_type = kTfLiteArenaRwPersistent;
+
+  return context->ResizeTensor(context, op_context.output,
+                               TfLiteIntArrayCopy(op_context.input->dims));
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
+  OpContext op_context(context, node);
+  if (op_data->dense_weights_initialized) {
+    return kTfLiteOk;
+  }
+
+  switch (op_context.input->type) {
+    case kTfLiteFloat32:
+      reference_ops::Densify(op_context.input->sparsity,
+                             GetTensorShape(op_context.input),
+                             GetTensorData<float>(op_context.input),
+                             GetTensorShape(op_context.output),
+                             GetTensorData<float>(op_context.output), context);
+      break;
+    case kTfLiteFloat16:
+      reference_ops::Densify(
+          op_context.input->sparsity, GetTensorShape(op_context.input),
+          GetTensorData<Eigen::half>(op_context.input),
+          GetTensorShape(op_context.output),
+          GetTensorData<Eigen::half>(op_context.output), context);
+      break;
+    case kTfLiteInt8:
+      reference_ops::Densify(op_context.input->sparsity,
+                             GetTensorShape(op_context.input),
+                             GetTensorData<int8_t>(op_context.input),
+                             GetTensorShape(op_context.output),
+                             GetTensorData<int8_t>(op_context.output), context);
+      break;
+
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %d not supported.",
+                         op_context.input->type);
+      return kTfLiteError;
+  }
+
+  op_data->dense_weights_initialized = true;
+  return kTfLiteOk;
+}
+
+}  // namespace densify
+
+TfLiteRegistration* Register_DENSIFY() {
+  static TfLiteRegistration r = {densify::Init, densify::Free, densify::Prepare,
+                                 densify::Eval};
+  return &r;
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/densify_test.cc

@@ -0,0 +1,111 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+namespace ops {
+namespace builtin {
+
+TfLiteRegistration* Register_DENSIFY();
+
+}  // namespace builtin
+}  // namespace ops
+
+namespace {
+
+using ::testing::ElementsAreArray;
+
+template <typename T>
+class DensifyOpModel : public SingleOpModel {
+ public:
+  DensifyOpModel(const TensorData& input, const std::vector<T>& input_data,
+                 int version = 1) {
+    input_ = AddConstSparseInput(input, input_data);
+    output_ = AddOutput({input.type, input.shape});
+
+    SetBuiltinOp(BuiltinOperator_DENSIFY, BuiltinOptions_DensifyOptions,
+                 CreateDensifyOptions(builder_).Union());
+
+    resolver_ = std::make_unique<SingleOpResolver>(
+        BuiltinOperator_DENSIFY, ops::builtin::Register_DENSIFY(), version);
+
+    BuildInterpreter({input.shape}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false, /*allocate_and_delegate=*/true);
+  }
+
+  std::vector<T> GetInput() { return ExtractVector<T>(input_); }
+  std::vector<T> GetOutput() { return ExtractVector<T>(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+TEST(DensifyOpTest, Float) {
+  std::vector<float> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  std::vector<float> sparse_values = {6, 9, 8, 5, 7};
+  TensorData input = {};
+  input.type = TensorType_FLOAT32;
+  input.shape = {3, 4};
+  input.traversal_order = {0, 1};
+  input.format = {kTfLiteDimDense, kTfLiteDimSparseCSR};
+  DensifyOpModel<float> m(input, dense_values);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetInput(), ElementsAreArray(sparse_values));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(dense_values));
+}
+
+TEST(DensifyOpTest, Float3D) {
+  std::vector<float> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  std::vector<float> sparse_values = {6, 9, 8, 5, 7};
+  TensorData input = {};
+  input.type = TensorType_FLOAT32;
+  input.shape = {3, 2, 2};
+  input.traversal_order = {0, 1, 2};
+  input.format = {kTfLiteDimDense, kTfLiteDimDense, kTfLiteDimSparseCSR};
+  DensifyOpModel<float> m(input, dense_values);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetInput(), ElementsAreArray(sparse_values));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(dense_values));
+}
+
+TEST(DensifyOpTest, Int8) {
+  std::vector<int8_t> dense_values = {6, 0, 9, 8, 0, 0, 0, 0, 5, 0, 0, 7};
+  std::vector<int8_t> sparse_values = {6, 9, 8, 5, 7};
+  TensorData input = {};
+  input.type = TensorType_INT8;
+  input.shape = {3, 4};
+  input.traversal_order = {0, 1};
+  input.format = {kTfLiteDimDense, kTfLiteDimSparseCSR};
+  DensifyOpModel<int8_t> m(input, dense_values);
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetInput(), ElementsAreArray(sparse_values));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(dense_values));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/deprecated_backends.cc

@@ -0,0 +1,24 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+namespace tflite {
+
+// Include this target as a dependency in order to define this function for
+// CpuBackendContext. Its use is to control execution of deprecated paths
+// by providing a symbol definition for otherwise "weak" symbol
+// declarations in CpuBackendContext.
+extern bool UseGemmlowpOnX86() { return true; }
+
+}  // namespace tflite

tensorflow/lite/kernels/depth_to_space.cc

@@ -0,0 +1,165 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace depth_to_space {
+
+// This file has two implementation of DepthToSpace. Note that DepthToSpace only
+// works on 4D tensors.
+enum KernelType {
+  kReference,
+  kGenericOptimized,
+};
+
+constexpr int kInputTensor = 0;
+constexpr int kOutputTensor = 0;
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteDepthToSpaceParams*>(node->builtin_data);
+
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
+
+  auto data_type = output->type;
+  TF_LITE_ENSURE(context,
+                 data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8 ||
+                     data_type == kTfLiteInt8 || data_type == kTfLiteInt32 ||
+                     data_type == kTfLiteInt64);
+  TF_LITE_ENSURE_TYPES_EQ(context, input->type, output->type);
+
+  const int block_size = params->block_size;
+  TF_LITE_ENSURE(context, block_size > 0);
+  const int input_height = input->dims->data[1];
+  const int input_width = input->dims->data[2];
+  const int input_channels = input->dims->data[3];
+  int output_height = input_height * block_size;
+  int output_width = input_width * block_size;
+  int output_channels = input_channels / block_size / block_size;
+
+  TF_LITE_ENSURE_EQ(context, input_height, output_height / block_size);
+  TF_LITE_ENSURE_EQ(context, input_width, output_width / block_size);
+  TF_LITE_ENSURE_EQ(context, input_channels,
+                    output_channels * block_size * block_size);
+
+  TfLiteIntArray* output_size = TfLiteIntArrayCreate(4);
+  output_size->data[0] = input->dims->data[0];
+  output_size->data[1] = output_height;
+  output_size->data[2] = output_width;
+  output_size->data[3] = output_channels;
+
+  return context->ResizeTensor(context, output, output_size);
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteDepthToSpaceParams*>(node->builtin_data);
+
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+#define TF_LITE_DEPTH_TO_SPACE(type, scalar)                               \
+  tflite::DepthToSpaceParams op_params;                                    \
+  op_params.block_size = params->block_size;                               \
+  type::DepthToSpace(op_params, GetTensorShape(input),                     \
+                     GetTensorData<scalar>(input), GetTensorShape(output), \
+                     GetTensorData<scalar>(output))
+  // Already know in/out types are same.
+  switch (TfLiteTypeGetSizeBits(input->type)) {
+    case 8:
+      if (kernel_type == kReference) {
+        TF_LITE_DEPTH_TO_SPACE(reference_ops, uint8_t);
+      } else {
+        TF_LITE_DEPTH_TO_SPACE(optimized_ops, uint8_t);
+      }
+      break;
+    case 16:
+      if (kernel_type == kReference) {
+        TF_LITE_DEPTH_TO_SPACE(reference_ops, uint16_t);
+      } else {
+        TF_LITE_DEPTH_TO_SPACE(optimized_ops, uint16_t);
+      }
+      break;
+    case 32:
+      if (kernel_type == kReference) {
+        TF_LITE_DEPTH_TO_SPACE(reference_ops, uint32_t);
+      } else {
+        TF_LITE_DEPTH_TO_SPACE(optimized_ops, uint32_t);
+      }
+      break;
+    case 64:
+      if (kernel_type == kReference) {
+        TF_LITE_DEPTH_TO_SPACE(reference_ops, uint64_t);
+      } else {
+        TF_LITE_DEPTH_TO_SPACE(optimized_ops, uint64_t);
+      }
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type '%s' not currently supported.",
+                         TfLiteTypeGetName(input->type));
+      return kTfLiteError;
+  }
+#undef TF_LITE_DEPTH_TO_SPACE
+
+  return kTfLiteOk;
+}
+
+}  // namespace depth_to_space
+
+TfLiteRegistration* Register_DEPTH_TO_SPACE_REF() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, depth_to_space::Prepare,
+      depth_to_space::Eval<depth_to_space::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTH_TO_SPACE_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      nullptr, nullptr, depth_to_space::Prepare,
+      depth_to_space::Eval<depth_to_space::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTH_TO_SPACE() {
+  return Register_DEPTH_TO_SPACE_GENERIC_OPT();
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/depth_to_space_test.cc

@@ -0,0 +1,116 @@
+/* Copyright 2017 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 <stdint.h>
+
+#include <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+
+class DepthToSpaceOpModel : public SingleOpModel {
+ public:
+  DepthToSpaceOpModel(const TensorData& tensor_data, int block_size) {
+    input_ = AddInput(tensor_data);
+    output_ = AddOutput(tensor_data);
+    SetBuiltinOp(BuiltinOperator_DEPTH_TO_SPACE,
+                 BuiltinOptions_DepthToSpaceOptions,
+                 CreateDepthToSpaceOptions(builder_, block_size).Union());
+    BuildInterpreter({GetShape(input_)});
+  }
+
+  template <typename T>
+  void SetInput(std::initializer_list<T> data) {
+    PopulateTensor<T>(input_, data);
+  }
+  template <typename T>
+  std::vector<T> GetOutput() {
+    return ExtractVector<T>(output_);
+  }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+
+ private:
+  int input_;
+  int output_;
+};
+
+#if GTEST_HAS_DEATH_TEST
+TEST(DepthToSpaceOpModel, BadBlockSize) {
+  EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 4),
+               "Cannot allocate tensors");
+}
+
+TEST(DepthToSpaceOpModel, NoBlockSize) {
+  EXPECT_DEATH(DepthToSpaceOpModel({TensorType_FLOAT32, {1, 1, 1, 4}}, 0),
+               "Cannot allocate tensors");
+}
+#endif
+
+TEST(DepthToSpaceOpModel, Float32) {
+  DepthToSpaceOpModel m({TensorType_FLOAT32, {1, 1, 1, 4}}, 2);
+  m.SetInput<float>({1.4, 2.3, 3.2, 4.1});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput<float>(),
+              Pointwise(FloatingPointEq(), {1.4, 2.3, 3.2, 4.1}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 2, 2, 1));
+}
+
+TEST(DepthToSpaceOpModel, Uint8) {
+  DepthToSpaceOpModel m({TensorType_UINT8, {1, 1, 2, 4}}, 2);
+  m.SetInput<uint8_t>({1, 2, 3, 4, 5, 6, 7, 8});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput<uint8_t>(),
+              ElementsAreArray({1, 2, 5, 6, 3, 4, 7, 8}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 2, 4, 1));
+}
+
+TEST(DepthToSpaceOpModel, int8) {
+  DepthToSpaceOpModel m({TensorType_INT8, {1, 2, 1, 4}}, 2);
+  m.SetInput<int8_t>({1, 2, 3, 4, 5, 6, 7, 8});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput<int8_t>(),
+              ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 4, 2, 1));
+}
+
+TEST(DepthToSpaceOpModel, Int32) {
+  DepthToSpaceOpModel m({TensorType_INT32, {1, 2, 2, 4}}, 2);
+  m.SetInput<int32_t>({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput<int32_t>(),
+              ElementsAreArray(
+                  {1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 4, 4, 1));
+}
+
+TEST(DepthToSpaceOpModel, Int64) {
+  DepthToSpaceOpModel m({TensorType_INT64, {1, 1, 1, 1}}, 1);
+  m.SetInput<int64_t>({4});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput<int64_t>(), ElementsAreArray({4}));
+  EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 1, 1, 1));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/depthwise_conv.cc

@@ -0,0 +1,719 @@
+/* Copyright 2017 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/internal/optimized/integer_ops/depthwise_conv.h"
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_multithread.h"
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_float.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/depthwise_conv.h"
+#include "tensorflow/lite/kernels/internal/tensor.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/padding.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace depthwise_conv {
+constexpr int kInputTensor = 0;
+constexpr int kFilterTensor = 1;
+constexpr int kBiasTensor = 2;
+constexpr int kOutputTensor = 0;
+
+// This file has three implementation of DepthwiseConv.
+enum KernelType {
+  kReference,
+  kGenericOptimized,  // Neon-free
+  kNeonOptimized,
+};
+
+const int kTensorNotAllocated = -1;
+
+struct OpData {
+  TfLitePaddingValues padding;
+  // The scaling factor from input to output (aka the 'real multiplier') can
+  // be represented as a fixed point multiplier plus a left shift.
+  int32_t output_multiplier;
+  int output_shift;
+  // The range of the fused activation layer. For example for kNone and
+  // uint8_t these would be 0 and 255.
+  int32_t output_activation_min;
+  int32_t output_activation_max;
+
+  // Per channel output multiplier and shift.
+  std::vector<int32_t> per_channel_output_multiplier;
+  std::vector<int> per_channel_output_shift;
+
+  // Hybrid per channel temporary tensors.
+  int input_quantized_id = kTensorNotAllocated;
+  int scaling_factors_id = kTensorNotAllocated;
+  int input_offset_id = kTensorNotAllocated;
+  int32_t input_quantized_index;
+  int32_t scaling_factors_index;
+  int32_t input_offset_index;
+};
+
+void* Init(TfLiteContext* context, const char* buffer, size_t length) {
+  // This is a builtin op, so we don't use the contents in 'buffer', if any.
+  // Instead, we allocate a new object to carry information from Prepare() to
+  // Eval().
+  return new OpData;
+}
+
+void Free(TfLiteContext* context, void* buffer) {
+  delete reinterpret_cast<OpData*>(buffer);
+}
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteDepthwiseConvParams*>(node->builtin_data);
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  bool has_bias = NumInputs(node) == 3;
+
+  TF_LITE_ENSURE(context, has_bias || NumInputs(node) == 2);
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFilterTensor, &filter));
+  const TfLiteTensor* bias = nullptr;
+
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+
+  TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(filter), 4);
+  TF_LITE_ENSURE(context, params->dilation_height_factor > 0);
+  TF_LITE_ENSURE(context, params->dilation_width_factor > 0);
+
+  // Validate the stride values
+  TF_LITE_ENSURE(context, params->stride_width <= 0xFFFF);
+  TF_LITE_ENSURE(context, params->stride_width > 0);
+  TF_LITE_ENSURE(context, params->stride_height <= 0xFFFF);
+  TF_LITE_ENSURE(context, params->stride_height > 0);
+
+  const TfLiteType data_type = input->type;
+
+  const TfLiteType filter_type = filter->type;
+  const bool is_hybrid =
+      data_type == kTfLiteFloat32 && filter_type == kTfLiteInt8;
+  TF_LITE_ENSURE(context,
+                 data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8 ||
+                     data_type == kTfLiteInt8 || data_type == kTfLiteInt16);
+  TF_LITE_ENSURE_TYPES_EQ(context, output->type, data_type);
+  if (!is_hybrid) {
+    TF_LITE_ENSURE(context, filter->type == data_type ||
+                                data_type == kTfLiteInt16 ||
+                                filter->type == kTfLiteInt4);
+  }
+  if (filter->type == kTfLiteInt4) {
+    int filter_num_elements = 0;
+    TF_LITE_ENSURE_MSG(
+        context, CheckedNumElements(filter, filter_num_elements) == kTfLiteOk,
+        "%s", "DepthwiseConv int4 filter has too many elements.");
+  }
+
+  if (data_type == kTfLiteInt16) {
+    TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0);
+    TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0);
+  }
+
+  // Filter in DepthwiseConv is expected to be [1, H, W, O].
+  TF_LITE_ENSURE_EQ(context, SizeOfDimension(filter, 0), 1);
+
+  if (has_bias) {
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBiasTensor, &bias));
+    if (data_type == kTfLiteUInt8 || data_type == kTfLiteInt8) {
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt32);
+      TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
+    } else if (data_type == kTfLiteInt16) {
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, kTfLiteInt64);
+      TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0);
+    } else {
+      TF_LITE_ENSURE_TYPES_EQ(context, bias->type, data_type);
+    }
+    TF_LITE_ENSURE_EQ(context, NumDimensions(bias), 1);
+    TF_LITE_ENSURE_EQ(context, SizeOfDimension(filter, 3),
+                      SizeOfDimension(bias, 0));
+  }
+
+  int channels_out = SizeOfDimension(filter, 3);
+  int width = SizeOfDimension(input, 2);
+  int height = SizeOfDimension(input, 1);
+  int filter_width = SizeOfDimension(filter, 2);
+  int filter_height = SizeOfDimension(filter, 1);
+  int batches = SizeOfDimension(input, 0);
+
+  // Matching GetWindowedOutputSize in TensorFlow.
+  auto padding = params->padding;
+  int out_width, out_height;
+
+  data->padding = ComputePaddingHeightWidth(
+      params->stride_height, params->stride_width,
+      params->dilation_height_factor, params->dilation_width_factor, height,
+      width, filter_height, filter_width, padding, &out_height, &out_width);
+
+  // Note that quantized inference requires that all tensors have their
+  // parameters set. This is usually done during quantized training or
+  // calibration.
+  if (data_type != kTfLiteFloat32) {
+    TF_LITE_ENSURE_EQ(context, filter->quantization.type,
+                      kTfLiteAffineQuantization);
+    TF_LITE_ENSURE(context, filter->quantization.type != kTfLiteNoQuantization);
+    const auto* affine_quantization =
+        reinterpret_cast<TfLiteAffineQuantization*>(
+            filter->quantization.params);
+    TF_LITE_ENSURE(context, affine_quantization);
+    TF_LITE_ENSURE(context, affine_quantization->scale);
+    TF_LITE_ENSURE(context, (affine_quantization->scale->size == 1 ||
+                             affine_quantization->scale->size == channels_out));
+
+    data->per_channel_output_multiplier.resize(channels_out);
+    data->per_channel_output_shift.resize(channels_out);
+    TF_LITE_ENSURE_STATUS(tflite::PopulateConvolutionQuantizationParams(
+        context, input, filter, bias, output, params->activation,
+        &data->output_multiplier, &data->output_shift,
+        &data->output_activation_min, &data->output_activation_max,
+        data->per_channel_output_multiplier.data(),
+        data->per_channel_output_shift.data(), channels_out));
+  }
+
+  if (is_hybrid) {
+    TF_LITE_ENSURE(context, filter->quantization.type != kTfLiteNoQuantization);
+    const auto* affine_quantization =
+        reinterpret_cast<TfLiteAffineQuantization*>(
+            filter->quantization.params);
+    TF_LITE_ENSURE(context, affine_quantization);
+    TF_LITE_ENSURE(context, affine_quantization->scale);
+    TF_LITE_ENSURE_EQ(
+        context, affine_quantization->scale->size,
+        filter->dims->data[affine_quantization->quantized_dimension]);
+    // Eval uses NumElements(input) in int arithmetic for per-batch
+    // quantization, so reject oversized shapes during Prepare.
+    int input_num_elements = 0;
+    TF_LITE_ENSURE_MSG(
+        context, CheckedNumElements(input, input_num_elements) == kTfLiteOk,
+        "%s", "DepthwiseConv hybrid input has too many elements.");
+
+    int temporaries_count = 0;
+    data->input_quantized_index = temporaries_count;
+    if (data->input_quantized_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->input_quantized_id));
+    }
+    ++temporaries_count;
+    data->scaling_factors_index = temporaries_count;
+    if (data->scaling_factors_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->scaling_factors_id));
+    }
+    ++temporaries_count;
+    data->input_offset_index = temporaries_count;
+    if (data->input_offset_id == kTensorNotAllocated) {
+      TF_LITE_ENSURE_OK(
+          context, context->AddTensors(context, 1, &data->input_offset_id));
+    }
+    ++temporaries_count;
+
+    TfLiteIntArrayFree(node->temporaries);
+    node->temporaries = TfLiteIntArrayCreate(temporaries_count);
+
+    node->temporaries->data[data->input_quantized_index] =
+        data->input_quantized_id;
+    TfLiteTensor* input_quantized;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, data->input_quantized_index,
+                                  &input_quantized));
+    input_quantized->type = kTfLiteInt8;
+    input_quantized->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) {
+      TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims);
+      TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized,
+                                                       input_quantized_size));
+    }
+    node->temporaries->data[data->scaling_factors_index] =
+        data->scaling_factors_id;
+    TfLiteTensor* scaling_factors;
+    TF_LITE_ENSURE_OK(
+        context, GetTemporarySafe(context, node, data->scaling_factors_index,
+                                  &scaling_factors));
+    scaling_factors->type = kTfLiteFloat32;
+    scaling_factors->allocation_type = kTfLiteArenaRw;
+    const int batch_size = SizeOfDimension(input, 0);
+    int scaling_dims[1] = {batch_size};
+    if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) {
+      std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>
+          scaling_factors_size(TfLiteIntArrayCreate(1), TfLiteIntArrayFree);
+      scaling_factors_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context,
+                        context->ResizeTensor(context, scaling_factors,
+                                              scaling_factors_size.release()));
+    }
+    node->temporaries->data[data->input_offset_index] = data->input_offset_id;
+    TfLiteTensor* input_offsets;
+    TF_LITE_ENSURE_OK(context,
+                      GetTemporarySafe(context, node, data->input_offset_index,
+                                       &input_offsets));
+    input_offsets->type = kTfLiteInt32;
+    input_offsets->allocation_type = kTfLiteArenaRw;
+    if (!TfLiteIntArrayEqualsArray(input_offsets->dims, 1, scaling_dims)) {
+      std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)>
+          input_offsets_size(TfLiteIntArrayCreate(1), TfLiteIntArrayFree);
+      input_offsets_size->data[0] = batch_size;
+      TF_LITE_ENSURE_OK(context,
+                        context->ResizeTensor(context, input_offsets,
+                                              input_offsets_size.release()));
+    }
+  }
+
+  std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> outputSize(
+      TfLiteIntArrayCreate(4), TfLiteIntArrayFree);
+  outputSize->data[0] = batches;
+  outputSize->data[1] = out_height;
+  outputSize->data[2] = out_width;
+  outputSize->data[3] = channels_out;
+  return context->ResizeTensor(context, output, outputSize.release());
+}
+
+TfLiteStatus ComputeDepthMultiplier(TfLiteContext* context,
+                                    const TfLiteTensor* input,
+                                    const TfLiteTensor* filter,
+                                    int16* depth_multiplier) {
+  int num_filter_channels = SizeOfDimension(filter, 3);
+  int num_input_channels = SizeOfDimension(input, 3);
+  TF_LITE_ENSURE(context, num_input_channels != 0);
+  TF_LITE_ENSURE_EQ(context, num_filter_channels % num_input_channels, 0);
+  *depth_multiplier = num_filter_channels / num_input_channels;
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node,
+                       TfLiteDepthwiseConvParams* params, OpData* data,
+                       const TfLiteTensor* input, const TfLiteTensor* filter,
+                       const TfLiteTensor* bias, TfLiteTensor* output) {
+  float output_activation_min, output_activation_max;
+  CalculateActivationRange(params->activation, &output_activation_min,
+                           &output_activation_max);
+
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
+  op_params.float_activation_min = output_activation_min;
+  op_params.float_activation_max = output_activation_max;
+  TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter,
+                                               &op_params.depth_multiplier));
+  if (kernel_type == kReference) {
+    reference_ops::DepthwiseConv(
+        op_params, GetTensorShape(input), GetTensorData<float>(input),
+        GetTensorShape(filter), GetTensorData<float>(filter),
+        GetTensorShape(bias), GetTensorData<float>(bias),
+        GetTensorShape(output), GetTensorData<float>(output));
+  } else {
+    optimized_ops::DepthwiseConv<float, float>(
+        op_params, GetTensorShape(input), GetTensorData<float>(input),
+        GetTensorShape(filter), GetTensorData<float>(filter),
+        GetTensorShape(bias), GetTensorData<float>(bias),
+        GetTensorShape(output), GetTensorData<float>(output),
+        CpuBackendContext::GetFromContext(context));
+  }
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node,
+                           TfLiteDepthwiseConvParams* params, OpData* data,
+                           const TfLiteTensor* input,
+                           const TfLiteTensor* filter, const TfLiteTensor* bias,
+                           TfLiteTensor* output) {
+  auto input_offset = -input->params.zero_point;
+  auto filter_offset = -filter->params.zero_point;
+  auto output_offset = output->params.zero_point;
+
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
+  op_params.input_offset = input_offset;
+  op_params.weights_offset = filter_offset;
+  op_params.output_offset = output_offset;
+  op_params.output_multiplier = data->output_multiplier;
+  op_params.output_shift = -data->output_shift;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+  TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter,
+                                               &op_params.depth_multiplier));
+  if (kernel_type == kReference) {
+    reference_ops::DepthwiseConv(
+        op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),
+        GetTensorShape(filter), GetTensorData<uint8_t>(filter),
+        GetTensorShape(bias), GetTensorData<int32_t>(bias),
+        GetTensorShape(output), GetTensorData<uint8_t>(output));
+  } else {
+    optimized_ops::DepthwiseConv<uint8, int32>(
+        op_params, GetTensorShape(input), GetTensorData<uint8_t>(input),
+        GetTensorShape(filter), GetTensorData<uint8_t>(filter),
+        GetTensorShape(bias), GetTensorData<int32_t>(bias),
+        GetTensorShape(output), GetTensorData<uint8_t>(output),
+        CpuBackendContext::GetFromContext(context));
+  }
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node,
+                                     TfLiteDepthwiseConvParams* params,
+                                     OpData* data, const TfLiteTensor* input,
+                                     const TfLiteTensor* filter,
+                                     const TfLiteTensor* bias,
+                                     TfLiteTensor* output) {
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
+  op_params.input_offset = -input->params.zero_point;
+  op_params.weights_offset = 0;
+  op_params.output_offset = output->params.zero_point;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+  TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter,
+                                               &op_params.depth_multiplier));
+
+  KernelType effective_kernel_type = kernel_type;
+
+  const int8_t* filter_data;
+  const size_t bytes_unpacked = filter->bytes * 2;
+  auto unpacked_filter_data = std::make_unique<int8_t[]>(bytes_unpacked);
+
+  if (filter->type == kTfLiteInt4) {
+    tflite::tensor_utils::UnpackPackedIntToInt8(
+        GetTensorData<int8_t>(filter), GetTensorShape(filter).FlatSize(),
+        /*bit_width=*/4, unpacked_filter_data.get());
+    filter_data = unpacked_filter_data.get();
+  } else {
+    filter_data = GetTensorData<int8>(filter);
+  }
+
+  switch (effective_kernel_type) {
+    case kReference: {
+      switch (filter->type) {
+        case kTfLiteInt4:
+        case kTfLiteInt8: {
+          reference_integer_ops::DepthwiseConvPerChannel(
+              op_params, data->per_channel_output_multiplier.data(),
+              data->per_channel_output_shift.data(), GetTensorShape(input),
+              GetTensorData<int8>(input), GetTensorShape(filter), filter_data,
+              GetTensorShape(bias), GetTensorData<int32>(bias),
+              GetTensorShape(output), GetTensorData<int8>(output));
+          break;
+        }
+        default: {
+          TF_LITE_KERNEL_LOG(context,
+                             "Weight type %s (%d) not supported for filter.",
+                             TfLiteTypeGetName(filter->type), filter->type);
+          break;
+        }
+      }
+      break;
+    }
+    case kGenericOptimized:
+    case kNeonOptimized: {
+      switch (filter->type) {
+        case kTfLiteInt4:
+        case kTfLiteInt8: {
+          optimized_integer_ops::DepthwiseConvPerChannel(
+              op_params, data->per_channel_output_multiplier.data(),
+              data->per_channel_output_shift.data(), GetTensorShape(input),
+              GetTensorData<int8>(input), GetTensorShape(filter), filter_data,
+              GetTensorShape(bias), GetTensorData<int32>(bias),
+              GetTensorShape(output), GetTensorData<int8>(output),
+              CpuBackendContext::GetFromContext(context));
+          break;
+        }
+        default: {
+          TF_LITE_KERNEL_LOG(context,
+                             "Weight type %s (%d) not supported for filter.",
+                             TfLiteTypeGetName(filter->type), filter->type);
+          break;
+        }
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus EvalQuantizedPerChannel16x8(
+    const TfLiteDepthwiseConvParams* params, const OpData* data,
+    const TfLiteTensor* input, const TfLiteTensor* filter,
+    const TfLiteTensor* bias, TfLiteTensor* output) {
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
+  op_params.depth_multiplier = params->depth_multiplier;
+  op_params.weights_offset = 0;
+  op_params.quantized_activation_min = data->output_activation_min;
+  op_params.quantized_activation_max = data->output_activation_max;
+
+  reference_integer_ops::DepthwiseConvPerChannel(
+      op_params, data->per_channel_output_multiplier.data(),
+      data->per_channel_output_shift.data(), GetTensorShape(input),
+      GetTensorData<int16>(input), GetTensorShape(filter),
+      GetTensorData<int8>(filter), GetTensorShape(bias),
+      GetTensorData<std::int64_t>(bias), GetTensorShape(output),
+      GetTensorData<int16>(output));
+
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type>
+TfLiteStatus EvalHybridPerChannel(TfLiteContext* context, TfLiteNode* node,
+                                  TfLiteDepthwiseConvParams* params,
+                                  OpData* data, const TfLiteTensor* input,
+                                  const TfLiteTensor* filter,
+                                  const TfLiteTensor* bias,
+                                  TfLiteTensor* output) {
+  float output_activation_min, output_activation_max;
+  CalculateActivationRange(params->activation, &output_activation_min,
+                           &output_activation_max);
+  const int batch_size = SizeOfDimension(input, 0);
+  TF_LITE_ENSURE(context, batch_size != 0);
+  const int input_size = NumElements(input) / batch_size;
+  TfLiteTensor* input_quantized;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->input_quantized_index,
+                                     &input_quantized));
+  int8_t* quantized_input_ptr_batch = input_quantized->data.int8;
+  TfLiteTensor* scaling_factors_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->scaling_factors_index,
+                                     &scaling_factors_tensor));
+  float* scaling_factors_ptr = GetTensorData<float>(scaling_factors_tensor);
+  TfLiteTensor* input_offset_tensor;
+  TF_LITE_ENSURE_OK(context,
+                    GetTemporarySafe(context, node, data->input_offset_index,
+                                     &input_offset_tensor));
+  int32_t* input_offset_ptr = GetTensorData<int32_t>(input_offset_tensor);
+
+  for (int b = 0; b < batch_size; ++b) {
+    const int offset = b * input_size;
+    tensor_utils::AsymmetricQuantizeFloats(
+        GetTensorData<float>(input) + offset, input_size,
+        quantized_input_ptr_batch + offset, &scaling_factors_ptr[b],
+        &input_offset_ptr[b]);
+  }
+
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = data->padding.width;
+  op_params.padding_values.height = data->padding.height;
+  op_params.stride_width = params->stride_width;
+  op_params.stride_height = params->stride_height;
+  op_params.dilation_width_factor = params->dilation_width_factor;
+  op_params.dilation_height_factor = params->dilation_height_factor;
+  op_params.depth_multiplier = params->depth_multiplier;
+
+  op_params.weights_offset = 0;
+  op_params.float_activation_min = output_activation_min;
+  op_params.float_activation_max = output_activation_max;
+  TF_LITE_ENSURE(context, filter->quantization.type != kTfLiteNoQuantization);
+  const auto* affine_quantization =
+      reinterpret_cast<TfLiteAffineQuantization*>(filter->quantization.params);
+  if (kernel_type == kReference) {
+    reference_integer_ops::DepthwiseConvHybridPerChannel(
+        op_params, scaling_factors_ptr, GetTensorShape(input),
+        quantized_input_ptr_batch, GetTensorShape(filter),
+        GetTensorData<int8>(filter), GetTensorShape(bias),
+        GetTensorData<float>(bias), GetTensorShape(output),
+        GetTensorData<float>(output), affine_quantization->scale->data,
+        input_offset_ptr);
+  } else {
+    optimized_integer_ops::DepthwiseConvHybridPerChannel(
+        op_params, scaling_factors_ptr, GetTensorShape(input),
+        quantized_input_ptr_batch, GetTensorShape(filter),
+        GetTensorData<int8>(filter), GetTensorShape(bias),
+        GetTensorData<float>(bias), GetTensorShape(output),
+        GetTensorData<float>(output), affine_quantization->scale->data,
+        input_offset_ptr, CpuBackendContext::GetFromContext(context));
+  }
+
+  return kTfLiteOk;
+}
+
+template <KernelType kernel_type, TfLiteType input_type>
+TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) {
+  auto* params =
+      reinterpret_cast<TfLiteDepthwiseConvParams*>(node->builtin_data);
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(context,
+                    GetOutputSafe(context, node, kOutputTensor, &output));
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+  const TfLiteTensor* filter;
+  TF_LITE_ENSURE_OK(context,
+                    GetInputSafe(context, node, kFilterTensor, &filter));
+  const TfLiteTensor* bias =
+      (NumInputs(node) == 3) ? GetInput(context, node, kBiasTensor) : nullptr;
+  TFLITE_DCHECK_EQ(input_type, input->type);
+
+  switch (input_type) {  // Already know in/out types are same.
+    case kTfLiteFloat32:
+      if (filter->type == kTfLiteFloat32) {
+        return EvalFloat<kernel_type>(context, node, params, data, input,
+                                      filter, bias, output);
+      } else if (filter->type == kTfLiteInt8) {
+        return EvalHybridPerChannel<kernel_type>(context, node, params, data,
+                                                 input, filter, bias, output);
+      } else {
+        TF_LITE_KERNEL_LOG(
+            context, "Type %s with filter type %s not currently supported.",
+            TfLiteTypeGetName(input->type), TfLiteTypeGetName(filter->type));
+        return kTfLiteError;
+      }
+      break;
+    case kTfLiteUInt8:
+      return EvalQuantized<kernel_type>(context, node, params, data, input,
+                                        filter, bias, output);
+      break;
+    case kTfLiteInt8:
+      return EvalQuantizedPerChannel<kernel_type>(context, node, params, data,
+                                                  input, filter, bias, output);
+      break;
+    case kTfLiteInt16:
+      return EvalQuantizedPerChannel16x8(params, data, input, filter, bias,
+                                         output);
+      break;
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.",
+                         input->type);
+      return kTfLiteError;
+  }
+}
+
+template <KernelType kernel_type>
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
+
+  switch (input->type) {  // Already know in/out types are same.
+    case kTfLiteFloat32:
+      return EvalImpl<kernel_type, kTfLiteFloat32>(context, node);
+    case kTfLiteUInt8:
+      return EvalImpl<kernel_type, kTfLiteUInt8>(context, node);
+    case kTfLiteInt8:
+      return EvalImpl<kernel_type, kTfLiteInt8>(context, node);
+    case kTfLiteInt16:
+      return EvalImpl<kernel_type, kTfLiteInt16>(context, node);
+    default:
+      TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.",
+                         input->type);
+      return kTfLiteError;
+  }
+}
+
+}  // namespace depthwise_conv
+
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_REF() {
+  static TfLiteRegistration r = {
+      depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare,
+      depthwise_conv::Eval<depthwise_conv::kReference>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT() {
+  static TfLiteRegistration r = {
+      depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare,
+      depthwise_conv::Eval<depthwise_conv::kGenericOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_NEON_OPT() {
+  static TfLiteRegistration r = {
+      depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare,
+      depthwise_conv::Eval<depthwise_conv::kNeonOptimized>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_NEON_OPT_UINT8() {
+  static TfLiteRegistration r = {
+      depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare,
+      depthwise_conv::EvalImpl<depthwise_conv::kNeonOptimized, kTfLiteUInt8>};
+  return &r;
+}
+
+TfLiteRegistration* Register_DEPTHWISE_CONV_2D() {
+#ifdef USE_NEON
+  return Register_DEPTHWISE_CONVOLUTION_NEON_OPT();
+#else
+  return Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT();
+#endif
+}
+
+// Warning: Clients using this variant are responsible for ensuring that their
+// models only need the UINT8 type. TFLite's op registration mechanism doesn't
+// yet allow for more nuanced registration mechanisms.
+TfLiteRegistration* Register_DEPTHWISE_CONV_2D_UINT8() {
+#ifdef USE_NEON
+  return Register_DEPTHWISE_CONVOLUTION_NEON_OPT_UINT8();
+#else
+  return Register_DEPTHWISE_CONV_2D();
+#endif
+}
+
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite

tensorflow/lite/kernels/depthwise_conv_hybrid_test.cc

@@ -0,0 +1,463 @@
+/* Copyright 2017 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 <stddef.h>
+
+#include <cstdint>
+#include <initializer_list>
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/memory/memory.h"
+#include "flatbuffers/flatbuffers.h"  // from @flatbuffers
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
+
+namespace tflite {
+
+namespace ops {
+namespace builtin {
+
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_REF();
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT();
+TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_NEON_OPT();
+
+}  // namespace builtin
+}  // namespace ops
+
+namespace {
+
+using ::testing::ElementsAreArray;
+
+class BaseDepthwiseConvolutionOpModel : public SingleOpModel {
+ public:
+  BaseDepthwiseConvolutionOpModel(
+      TfLiteRegistration* registration, const TensorData& input,
+      const TensorData& filter, const TensorData& output, Padding padding_type,
+      int dilation_factor = 1, int stride_width = 1, int stride_height = 1,
+      ActivationFunctionType fused_activation_function =
+          ActivationFunctionType_NONE) {
+    input_ = AddInput(input);
+    filter_ = AddInput(filter);
+
+    int bias_size = GetShape(filter_)[3];
+    if (input.type == TensorType_FLOAT32) {
+      bias_ = AddInput({TensorType_FLOAT32, {bias_size}});
+    } else {
+      // This is a quantized version. The scale of 'bias' depends on the scales
+      // of input and filter. Supposedly this is correctly set during quantized
+      // training.
+      if (filter.per_channel_quantization) {
+        // per channel quantization.
+        std::vector<float> bias_scale(
+            filter.per_channel_quantization_scales.size());
+        std::vector<int64_t> bias_zero_points(
+            filter.per_channel_quantization_scales.size());
+        for (size_t i = 0; i < filter.per_channel_quantization_scales.size();
+             ++i) {
+          bias_scale[i] =
+              input.scale * filter.per_channel_quantization_scales[i];
+          bias_zero_points[i] = 0;
+        }
+        TensorData bias{TensorType_INT32,
+                        {bias_size},
+                        /*min=*/0,
+                        /*max=*/0,
+                        /*scale=*/0,
+                        /*zero_point=*/0,
+                        true,
+                        /*per_channel_quantization_scales=*/bias_scale,
+                        /*per_channel_quantization_offsets=*/bias_zero_points,
+                        /*channel_index==*/0};
+        bias_ = AddInput(bias);
+      } else {
+        // per tensor quantization.
+        auto bias_scale = GetScale(input_) * GetScale(filter_);
+        TensorData bias{TensorType_INT32, {bias_size}, 0, 0, bias_scale};
+        bias_ = AddInput(bias);
+      }
+    }
+
+    output_ = AddOutput(output);
+
+    int input_depth = GetShape(input_)[3];
+    int output_depth = GetShape(filter_)[3];
+    int depth_mul = output_depth / input_depth;
+
+    SetBuiltinOp(
+        BuiltinOperator_DEPTHWISE_CONV_2D,
+        BuiltinOptions_DepthwiseConv2DOptions,
+        CreateDepthwiseConv2DOptions(
+            builder_, padding_type, stride_width, stride_height, depth_mul,
+            fused_activation_function, dilation_factor, dilation_factor)
+            .Union());
+
+    resolver_ = std::make_unique<SingleOpResolver>(
+        BuiltinOperator_DEPTHWISE_CONV_2D, registration);
+
+    BuildInterpreter({GetShape(input_), GetShape(filter_), GetShape(bias_)});
+  }
+
+ protected:
+  int input_;
+  int filter_;
+  int bias_;
+  int output_;
+};
+
+class PerChannelHybridDepthwiseConvolutionOpModel
+    : public BaseDepthwiseConvolutionOpModel {
+ public:
+  using BaseDepthwiseConvolutionOpModel::BaseDepthwiseConvolutionOpModel;
+
+  void SetInput(std::initializer_list<float> data) {
+    PopulateTensor(input_, data);
+  }
+
+  void SetFilter(std::initializer_list<float> data) {
+    PerChannelSymmetricQuantizeAndPopulate(filter_, data);
+  }
+
+  void SetBias(std::initializer_list<float> data) {
+    PopulateTensor(bias_, data);
+  }
+
+  void SetInput(const std::vector<float>& data) {
+    PopulateTensor(input_, data);
+  }
+
+  void SetFilter(const std::vector<float>& data) {
+    PerChannelSymmetricQuantizeAndPopulate(filter_, data);
+  }
+
+  void SetBias(const std::vector<float>& data) { PopulateTensor(bias_, data); }
+
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+};
+
+const auto kKernelMap = new std::map<string, TfLiteRegistration*>({
+    {"Reference", ops::builtin::Register_DEPTHWISE_CONVOLUTION_REF()},
+    {"GenericOptimized",
+     ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT()},
+    {"NeonOptimized", ops::builtin::Register_DEPTHWISE_CONVOLUTION_NEON_OPT()},
+});
+
+class PerChannelHybridDepthwiseConvolutionOptimizedOpTest
+    : public SingleOpTest {
+ protected:
+  const std::map<string, TfLiteRegistration*>& GetKernelMap() override {
+    return *kKernelMap;
+  }
+};
+
+class PerChannelHybridDepthwiseConvolutionOpTest : public SingleOpTest {
+ protected:
+  const std::map<string, TfLiteRegistration*>& GetKernelMap() override {
+    return *kKernelMap;
+  }
+};
+
+void RandomTest(int b, int h, int w, int c, int fs, bool padding, int sw) {
+  const float element_max = 1.0;
+  const int input_size = b * h * w * c;
+  const int filter_size = 1 * fs * fs * c;
+  const int bias_size = c;
+  std::vector<float> input_data(input_size);
+  std::vector<float> filter_data(filter_size);
+  std::vector<float> bias_data(bias_size);
+  for (int i = 0; i < input_size; ++i) {
+    input_data[i] = UniformRandomFloat(-element_max, element_max);
+  }
+  for (int i = 0; i < filter_size; ++i) {
+    filter_data[i] = UniformRandomFloat(-element_max, element_max);
+  }
+  for (int i = 0; i < bias_size; ++i) {
+    bias_data[i] = UniformRandomFloat(-element_max, element_max);
+  }
+  const TensorData input({TensorType_FLOAT32, {b, h, w, c}});
+  const TensorData output({TensorType_FLOAT32, {}});
+  std::vector<float> scales;
+  std::vector<int64_t> offsets;
+  for (int i = 0; i < c; i++) {
+    scales.push_back(1.0 / 127.0);
+    offsets.push_back(0.0);
+  }
+  const TensorData filter({TensorType_INT8,
+                           {1, fs, fs, c},
+                           0,
+                           0,
+                           0,
+                           0,
+                           /*per_channel_quantization=*/true,
+                           scales,
+                           offsets,
+                           3});
+  PerChannelHybridDepthwiseConvolutionOpModel hybrid_generic(
+      ops::builtin::Register_DEPTHWISE_CONVOLUTION_REF(), input, filter, output,
+      padding ? Padding_SAME : Padding_VALID,
+      /* dilation_factor = */ 1,
+      /* stride_width = */ sw,
+      /* stride_height = */ sw);
+  hybrid_generic.SetInput(input_data);
+  hybrid_generic.SetFilter(filter_data);
+  hybrid_generic.SetBias(bias_data);
+  ASSERT_EQ(hybrid_generic.Invoke(), kTfLiteOk);
+  std::vector<float> hybrid_generic_output = hybrid_generic.GetOutput();
+  PerChannelHybridDepthwiseConvolutionOpModel hybrid_optimized(
+      ops::builtin::Register_DEPTHWISE_CONVOLUTION_NEON_OPT(), input, filter,
+      output, padding ? Padding_SAME : Padding_VALID,
+      /* dilation_factor = */ 1,
+      /* stride_width = */ sw,
+      /* stride_height = */ sw);
+  hybrid_optimized.SetInput(input_data);
+  hybrid_optimized.SetFilter(filter_data);
+  hybrid_optimized.SetBias(bias_data);
+  ASSERT_EQ(hybrid_optimized.Invoke(), kTfLiteOk);
+  std::vector<float> hybrid_optimized_output = hybrid_optimized.GetOutput();
+  EXPECT_THAT(hybrid_generic_output,
+              ElementsAreArray(ArrayFloatNear(hybrid_optimized_output)));
+}
+
+void RandomTest(int b, int w, int h, int c, int fs) {
+  RandomTest(b, w, h, c, fs, false, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest32) {
+  RandomTest(1, 10, 10, 8, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest64) {
+  RandomTest(1, 112, 112, 64, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest128) {
+  RandomTest(1, 56, 56, 128, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest256) {
+  RandomTest(1, 28, 28, 256, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest512) {
+  RandomTest(1, 14, 14, 512, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest, AccuracyTest1024) {
+  RandomTest(1, 3, 3, 1024, 3);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest32) {
+  RandomTest(1, 112, 112, 32, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest64) {
+  RandomTest(1, 112, 112, 64, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest128) {
+  RandomTest(1, 56, 56, 128, 3, true, 1);
+}
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest256) {
+  RandomTest(1, 28, 28, 256, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest512) {
+  RandomTest(1, 14, 14, 512, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddingTest1024) {
+  RandomTest(1, 3, 3, 1024, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       AccuracyPaddiacc_buffer_sizengTest4096) {
+  RandomTest(1, 3, 3, 4096, 3, true, 1);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest32) {
+  RandomTest(1, 112, 112, 32, 3, false, 2);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest64) {
+  RandomTest(1, 112, 112, 64, 3, false, 2);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest128) {
+  RandomTest(1, 56, 56, 128, 3, false, 2);
+}
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest256) {
+  RandomTest(1, 28, 28, 256, 3, false, 2);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest512) {
+  RandomTest(1, 14, 14, 512, 3, false, 2);
+}
+
+TEST_F(PerChannelHybridDepthwiseConvolutionOptimizedOpTest,
+       Accuracy2x2StrideTest1024) {
+  RandomTest(1, 3, 3, 1024, 3, false, 1);
+}
+
+TEST_P(PerChannelHybridDepthwiseConvolutionOpTest, SimpleTest) {
+  PerChannelHybridDepthwiseConvolutionOpModel m(
+      GetRegistration(), {TensorType_FLOAT32, {1, 2, 3, 2}},
+      {TensorType_INT8,
+       // [1 * 2 * 2 * 4] as [input_channel, y, x, output_channel]
+       {1, 2, 2, 4},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/{1, 2, 3, 4},
+       /*per_channel_quantization_offsets=*/{0, 0, 0, 0},
+       /*channel_index=*/3},
+      {TensorType_FLOAT32, {}}, Padding_VALID);
+  m.SetInput({
+      // [1 * 2 * 3 * 2] as [batch, y, x, input_channel]
+      3, 2,    // batch = 0, y = 0, x = 0
+      1, -1,   // batch = 0, y = 0, x = 1
+      -2, -3,  // batch = 0, y = 0, x = 2
+      4, 3,    // batch = 0, y = 1, x = 0
+      2, -2,   // batch = 0, y = 1, x = 1
+      -3, -4,  // batch = 0, y = 1, x = 2
+  });
+  m.SetFilter(
+      /*filter data*/
+      {
+          // [1 * 2 * 2 * 4] as [input_channel, y, x, output_channel]
+          // depth multiplier = 2
+          1, 2, 3, 4,  // y = 0, x = 0
+          3, 4, 5, 6,  // y = 0, x = 1
+          7, 8, 5, 6,  // y = 1, x = 0
+          3, 4, 1, 2,  // y = 1, x = 1
+      });
+  m.SetBias({3, -2, 4, 6});
+
+  // Invoke and verify output.
+  // output has dimension [1 * 1 * 2 * 4] as [batch, y, x, output_channel]
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(
+      m.GetOutput(),
+      ElementsAreArray(ArrayFloatNear(
+          {42.9373, 47.9451, 22.0706, 22.0627, 3, -4.00784, -29.1294, -54.1098},
+          0.16)));
+}
+
+TEST_P(PerChannelHybridDepthwiseConvolutionOpTest, Simple3x3FilterTest) {
+  PerChannelHybridDepthwiseConvolutionOpModel m(
+      GetRegistration(), {TensorType_FLOAT32, {1, 3, 3, 8}},
+      {TensorType_INT8,
+       // [1 * 3 * 3 * 8] as [input_channel, y, x, output_channel]
+       {1, 3, 3, 8},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/
+       {1, 2, 3, 4, 4, 3, 2, 1},
+       /*per_channel_quantization_offsets=*/{0, 0, 0, 0, 0, 0, 0, 0},
+       /*channel_index=*/3},
+      {TensorType_FLOAT32, {}}, Padding_VALID);
+  m.SetInput({// array of 9 x 8 => [1, 3, 3, 8]
+              1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1,
+              0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0,
+              1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1,
+              0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0});
+  m.SetFilter(
+      /*filter data*/
+      {// array of 9 x 8 => [1, 3, 3, 8]
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8,
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8,
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8});
+  m.SetBias({0, 0, 0, 0, 0, 0, 0, 0});
+
+  // Invoke and verify output.
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(
+                                 {9, 18, 0, 0, 36, 54, 0, 0}, 0.16)));
+}
+
+TEST_P(PerChannelHybridDepthwiseConvolutionOpTest,
+       Simple3x3FilterPaddingSameTest) {
+  PerChannelHybridDepthwiseConvolutionOpModel m(
+      GetRegistration(), {TensorType_FLOAT32, {1, 3, 3, 8}},
+      {TensorType_INT8,
+       // [1 * 3 * 3 * 8] as [input_channel, y, x, output_channel]
+       {1, 3, 3, 8},
+       0,
+       0,
+       0,
+       0,
+       /*per_channel_quantization=*/true,
+       /*per_channel_quantization_scales=*/
+       {1, 2, 3, 4, 4, 3, 2, 1},
+       /*per_channel_quantization_offsets=*/{0, 0, 0, 0, 0, 0, 0, 0},
+       /*channel_index=*/3},
+      {TensorType_FLOAT32, {}}, Padding_SAME);
+  m.SetInput({// array of 9 x 8 => [1, 3, 3, 8]
+              1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1,
+              0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0,
+              1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1,
+              0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0});
+  m.SetFilter(
+      /*filter data*/
+      {// array of 9 x 8 => [1, 3, 3, 8]
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8,
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8,
+       1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8});
+  m.SetBias({0, 0, 0, 0, 0, 0, 0, 0});
+
+  // Invoke and verify output.
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray(ArrayFloatNear(
+                  {
+                      // array of 9 x 8 => [1, 3, 3, 8]
+                      4,  8,  0, 0,  16, 24, 0,  0,  6,  12, 0,  0,  24, 36, 0,
+                      0,  4,  8, 0,  0,  16, 24, 0,  0,  6,  12, 0,  0,  24, 36,
+                      0,  0,  9, 18, 0,  0,  36, 54, 0,  0,  6,  12, 0,  0,  24,
+                      36, 0,  0, 4,  8,  0,  0,  16, 24, 0,  0,  6,  12, 0,  0,
+                      24, 36, 0, 0,  4,  8,  0,  0,  16, 24, 0,  0,
+                  },
+                  0.16)));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    PerChannelHybridDepthwiseConvolutionOpTest,
+    PerChannelHybridDepthwiseConvolutionOpTest,
+    ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
+
+}  // namespace
+}  // namespace tflite