chore: import upstream snapshot with attribution

tensorflow/lite/experimental/ml_adjacent/BUILD

@@ -0,0 +1,12 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "lib",
+    hdrs = ["lib.h"],
+)

tensorflow/lite/experimental/ml_adjacent/algo/BUILD

@@ -0,0 +1,208 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "crop",
+    srcs = ["crop.cc"],
+    hdrs = [
+        "crop.h",
+    ],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "crop_test",
+    srcs = ["crop_test.cc"],
+    deps = [
+        ":crop",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "resize",
+    srcs = ["resize.cc"],
+    hdrs = ["resize.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "resize_test",
+    srcs = ["resize_test.cc"],
+    deps = [
+        ":resize",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "rotate",
+    srcs = ["rotate.cc"],
+    hdrs = ["rotate.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "rotate_test",
+    srcs = ["rotate_test.cc"],
+    deps = [
+        ":rotate",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "flip_left_right",
+    srcs = ["flip_left_right.cc"],
+    hdrs = ["flip_left_right.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "flip_left_right_test",
+    srcs = ["flip_left_right_test.cc"],
+    deps = [
+        ":flip_left_right",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "flip_up_down",
+    srcs = ["flip_up_down.cc"],
+    hdrs = ["flip_up_down.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "flip_up_down_test",
+    srcs = ["flip_up_down_test.cc"],
+    deps = [
+        ":flip_up_down",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "per_image_standardization",
+    srcs = ["per_image_standardization.cc"],
+    hdrs = ["per_image_standardization.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "per_image_standardization_test",
+    srcs = ["per_image_standardization_test.cc"],
+    deps = [
+        ":per_image_standardization",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "rgb_to_grayscale",
+    srcs = ["rgb_to_grayscale.cc"],
+    hdrs = ["rgb_to_grayscale.h"],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "rgb_to_grayscale_test",
+    srcs = ["rgb_to_grayscale_test.cc"],
+    deps = [
+        ":rgb_to_grayscale",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "rgb_to_yuv",
+    srcs = ["rgb_to_yuv.cc"],
+    hdrs = [
+        "image_utils.h",
+        "rgb_to_yuv.h",
+    ],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "rgb_to_yuv_test",
+    srcs = ["rgb_to_yuv_test.cc"],
+    deps = [
+        ":rgb_to_yuv",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "yuv_to_rgb",
+    srcs = ["yuv_to_rgb.cc"],
+    hdrs = [
+        "image_utils.h",
+        "yuv_to_rgb.h",
+    ],
+    deps = [
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "yuv_to_rgb_test",
+    srcs = ["yuv_to_rgb_test.cc"],
+    deps = [
+        ":yuv_to_rgb",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/data:owning_vector_ref",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/ml_adjacent/algo/crop.cc

@@ -0,0 +1,150 @@
+/* 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/experimental/ml_adjacent/algo/crop.h"
+
+#include <cstring>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace crop {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+using ::ml_adj::data::TypeWidth;
+
+// Crops given input to the bounding box. Works on any datatype.
+// Output buffer must be already resized.
+inline void CropToBoundingBox(dim_t offset_height, dim_t offset_width,
+                              dim_t out_height, dim_t out_width,
+                              const DataRef* input, MutableDataRef* output) {
+  const dim_t in_height = input->Dims()[1];
+  const dim_t in_width = input->Dims()[2];
+  const dim_t num_channels = input->Dims()[3];
+  const ind_t chunk =
+      static_cast<ind_t>(TypeWidth(input->Type())) * num_channels;
+  const ind_t in_img_size = static_cast<ind_t>(in_height) * in_width;
+  const ind_t out_img_size = static_cast<ind_t>(out_height) * out_width;
+
+  for (ind_t b = 0; b < input->Dims()[0]; ++b) {
+    for (ind_t i = 0; i < out_height; ++i) {
+      const ind_t read_byte_ofs =
+          (in_img_size * b + (i + offset_height) * in_width + offset_width) *
+          chunk;
+
+      const void* read_start_addr =
+          reinterpret_cast<const char*>(input->Data()) + read_byte_ofs;
+
+      const ind_t write_byte_ofs = chunk * (out_img_size * b + i * out_width);
+
+      void* write_addr =
+          reinterpret_cast<char*>(output->Data()) + write_byte_ofs;
+
+      // Copy slice of each input row by row.
+      std::memcpy(write_addr, read_start_addr, chunk * out_width);
+    }
+  }
+}
+
+// Crop given input from the center. Works on any datatype.
+void ComputeCenterCrop(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 2);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+
+  const DataRef* frac = inputs[1];
+  const double frac_data = *reinterpret_cast<const double*>(frac->Data());
+  TFLITE_CHECK(frac_data > 0.0 && frac_data <= 1.0);
+
+  // Compute output height.
+  const dim_t in_height = img->Dims()[1];
+  const dim_t out_height_offset =
+      static_cast<dim_t>((in_height - in_height * frac_data) / 2.0);
+  const dim_t out_height = in_height - (2 * out_height_offset);
+
+  // Compute output width.
+  const dim_t in_width = img->Dims()[2];
+  const dim_t out_width_offset =
+      static_cast<dim_t>((in_width - in_width * frac_data) / 2.0);
+  const dim_t out_width = in_width - (2 * out_width_offset);
+
+  // Resize output buffer.
+  MutableDataRef* output = outputs[0];
+  output->Resize({img->Dims()[0], out_height, out_width, img->Dims()[3]});
+
+  CropToBoundingBox(out_height_offset, out_width_offset, out_height, out_width,
+                    img, output);
+}
+
+// Crops given input to the bounding box. Works on any datatype.
+void ComputeCropToBoundingBox(const InputPack& inputs,
+                              const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 5);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract inputs.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+
+  const DataRef* offset_height = inputs[1];
+  const dim_t offset_height_data =
+      *reinterpret_cast<const dim_t*>(offset_height->Data());
+
+  const DataRef* offset_width = inputs[2];
+  const dim_t offset_width_data =
+      *reinterpret_cast<const dim_t*>(offset_width->Data());
+
+  const DataRef* target_height = inputs[3];
+  const dim_t target_height_data =
+      *reinterpret_cast<const dim_t*>(target_height->Data());
+
+  const DataRef* target_width = inputs[4];
+  const dim_t target_width_data =
+      *reinterpret_cast<const dim_t*>(target_width->Data());
+
+  TFLITE_CHECK_LE(offset_height_data + target_height_data, img->Dims()[1]);
+  TFLITE_CHECK_LE(offset_width_data + target_width_data, img->Dims()[2]);
+
+  // Resize output buffer.
+  MutableDataRef* output = outputs[0];
+  output->Resize(
+      {img->Dims()[0], target_height_data, target_width_data, img->Dims()[3]});
+
+  CropToBoundingBox(offset_height_data, offset_width_data, target_height_data,
+                    target_width_data, img, output);
+}
+
+}  // namespace
+
+const Algo* Impl_CenterCrop() {
+  static constexpr Algo kCenterCrop = {&ComputeCenterCrop, nullptr};
+  return &kCenterCrop;
+}
+
+const Algo* Impl_CropToBoundingBox() {
+  static constexpr Algo kCropToBoundingBox = {&ComputeCropToBoundingBox,
+                                              nullptr};
+  return &kCropToBoundingBox;
+}
+
+}  // namespace crop
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/crop.h

@@ -0,0 +1,56 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_CROP_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_CROP_H_
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace crop {
+
+// Center Crop
+//
+// Inputs: [img: any, frac: scalar<double>]
+// Ouputs: [img: any]
+//
+// Crop the given image from center, scaling height and width dimensions by
+// decimal input. Let `offs = floor((d - d * frac) / 2)` where `d` is a given
+// dimension and `frac` is the second decimal scalar input to this Algo. The
+// bounding box for each dimension `d` then covers `[offs, d - offs)`. Mimics
+// semantic of `tf.image.central_crop`.
+//
+// https://www.tensorflow.org/api_docs/python/tf/image/central_crop
+
+const algo::Algo* Impl_CenterCrop();
+
+// Crop to Bounding Box
+//
+// Inputs: [img: any, offset_height: unsigned, offset_width: unsigned,
+//          target_height: unsigned, target_width: unsigned]
+// Ouputs: [img: any]
+//
+// Cuts a rectangular bounding box out of given image. The top-left corner of
+// the bounding box is at (`offset_height`, `offset_width`) in image, and the
+// lower-right corner is at (`offset_height` + `target_height`,
+// `offset_width` + `target_width`). Mimics semantic of
+// `tf.image.crop_to_bounding_box`.
+//
+// https://www.tensorflow.org/api_docs/python/tf/image/crop_to_bounding_box
+
+const algo::Algo* Impl_CropToBoundingBox();
+
+}  // namespace crop
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_CROP_H_

tensorflow/lite/experimental/ml_adjacent/algo/crop_test.cc

@@ -0,0 +1,283 @@
+/* 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/experimental/ml_adjacent/algo/crop.h"
+
+#include <cstring>
+#include <numeric>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+using ::testing::ElementsAreArray;
+
+namespace ml_adj {
+namespace crop {
+namespace {
+
+// Gets a flattened vector from the given 4d coordinates whose values
+// range from [0, prod(d1, d2, d3, d4)).
+std::vector<float> GetIotaVec(dim_t d1, dim_t d2, dim_t d3, dim_t d4) {
+  std::vector<float> res;
+  res.resize(d1 * d2 * d3 * d4);
+  std::iota(res.begin(), res.end(), 0);
+  return res;
+}
+
+//----------------------------------------------------------------------------//
+//                            CENTRAL CROP TESTS                              //
+//----------------------------------------------------------------------------//
+
+struct CropCenterTestParams {
+  std::vector<dim_t> img_dims;
+  std::vector<float> img_data;
+  double frac;
+  std::vector<float> expected_data;
+  std::vector<dim_t> expected_shape;
+};
+
+class CropCenterTest : public testing::TestWithParam<CropCenterTestParams> {};
+
+TEST_P(CropCenterTest, FloatPixelType) {
+  const CropCenterTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Frac input.
+  OwningVectorRef frac(etype_t::f64);
+  frac.Resize({1});
+  ASSERT_EQ(frac.Bytes(), sizeof(double));
+  std::memcpy(frac.Data(), &params.frac, frac.Bytes());
+
+  // Empty output.
+  OwningVectorRef output(etype_t::f32);
+
+  const Algo* center_crop = Impl_CenterCrop();
+  center_crop->process({&img, &frac}, {&output});
+
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  EXPECT_THAT(absl::MakeSpan(out_data, output.NumElements()),
+              ElementsAreArray(params.expected_data));
+}
+
+
+INSTANTIATE_TEST_SUITE_P(
+    CropTests, CropCenterTest,
+    testing::ValuesIn({
+        CropCenterTestParams{{1, 4, 4, 1},
+                             GetIotaVec(1, 4, 4, 1),
+                             0.5,
+                             {5, 6, 9, 10},
+                             {1, 2, 2, 1}},
+        CropCenterTestParams{{1, 5, 5, 1},
+                             GetIotaVec(1, 5, 5, 1),
+                             0.5,
+                             {6, 7, 8, 11, 12, 13, 16, 17, 18},
+                             {1, 3, 3, 1}},
+        CropCenterTestParams{{1, 3, 3, 1},
+                             GetIotaVec(1, 3, 3, 1),
+                             0.5,
+                             {0, 1, 2, 3, 4, 5, 6, 7, 8},
+                             {1, 3, 3, 1}},
+        CropCenterTestParams{{1, 5, 5, 1},
+                             GetIotaVec(1, 5, 5, 1),
+                             0.9,
+                             GetIotaVec(1, 5, 5, 1),
+                             {1, 5, 5, 1}},
+        CropCenterTestParams{
+            {1, 5, 5, 1}, GetIotaVec(1, 5, 5, 1), 0.2, {12}, {1, 1, 1, 1}},
+        CropCenterTestParams{{1, 2, 2, 2},
+                             GetIotaVec(1, 2, 2, 2),
+                             .7,
+                             {0, 1, 2, 3, 4, 5, 6, 7},
+                             {1, 2, 2, 2}},
+        CropCenterTestParams{
+            {1, 3, 3, 2}, GetIotaVec(1, 3, 3, 2), .1, {8, 9}, {1, 1, 1, 2}},
+        CropCenterTestParams{
+            {2, 3, 3, 1}, GetIotaVec(2, 3, 3, 1), .1, {4, 13}, {2, 1, 1, 1}},
+        CropCenterTestParams{{2, 3, 3, 2},
+                             GetIotaVec(2, 3, 3, 2),
+                             .1,
+                             {8, 9, 26, 27},
+                             {2, 1, 1, 2}},
+    }));
+
+//----------------------------------------------------------------------------//
+//                        CROP TO BOUNDING BOX TESTS                          //
+//----------------------------------------------------------------------------//
+
+struct CropToBoundingBoxTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  dim_t offset_height;
+  dim_t offset_width;
+  dim_t target_height;
+  dim_t target_width;
+  const std::vector<dim_t> expected_shape;
+  const std::vector<float> expected_data;
+};
+
+class CropToBoundingBoxTest
+    : public testing::TestWithParam<CropToBoundingBoxTestParams> {};
+
+TEST_P(CropToBoundingBoxTest, FloatPixelType) {
+  const CropToBoundingBoxTestParams& params = GetParam();
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Offset height input.
+  OwningVectorRef offset_height(etype_t::i32);
+  offset_height.Resize({1});
+  ASSERT_EQ(offset_height.Bytes(), sizeof(int));
+  std::memcpy(offset_height.Data(), &params.offset_height,
+              offset_height.Bytes());
+
+  // Offset width input.
+  OwningVectorRef offset_width(etype_t::i32);
+  offset_width.Resize({1});
+  ASSERT_EQ(offset_width.Bytes(), sizeof(int));
+  std::memcpy(offset_width.Data(), &params.offset_width, offset_width.Bytes());
+
+  // Target height input.
+  OwningVectorRef target_height(etype_t::i32);
+  target_height.Resize({1});
+  ASSERT_EQ(target_height.Bytes(), sizeof(int));
+  std::memcpy(target_height.Data(), &params.target_height,
+              target_height.Bytes());
+
+  // Target width input.
+  OwningVectorRef target_width(etype_t::i32);
+  target_width.Resize({1});
+  ASSERT_EQ(target_width.Bytes(), sizeof(int));
+  std::memcpy(target_width.Data(), &params.target_width, target_width.Bytes());
+
+  // Empty output.
+  OwningVectorRef output(etype_t::f32);
+  const Algo* crop_to_bounding_box = Impl_CropToBoundingBox();
+  crop_to_bounding_box->process(
+      {&img, &offset_height, &offset_width, &target_height, &target_width},
+      {&output});
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  EXPECT_THAT(absl::MakeSpan(out_data, output.NumElements()),
+              ElementsAreArray(params.expected_data));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CropTests, CropToBoundingBoxTest,
+    testing::ValuesIn({
+        // Top-left corner.
+        CropToBoundingBoxTestParams{{1, 5, 5, 1},
+                                    GetIotaVec(1, 5, 5, 1),
+                                    0,
+                                    0,
+                                    2,
+                                    2,
+                                    {1, 2, 2, 1},
+                                    {0, 1,  //
+                                     5, 6}},
+        // Bottom-right corner.
+        CropToBoundingBoxTestParams{{1, 5, 5, 1},
+                                    GetIotaVec(1, 5, 5, 1),
+                                    3,
+                                    3,
+                                    2,
+                                    2,
+                                    {1, 2, 2, 1},
+                                    {18, 19,  //
+                                     23, 24}},
+        // Top-right corner.
+        CropToBoundingBoxTestParams{{1, 5, 5, 1},
+                                    GetIotaVec(1, 5, 5, 1),
+                                    0,
+                                    3,
+                                    2,
+                                    2,
+                                    {1, 2, 2, 1},
+                                    {3, 4,  //
+                                     8, 9}},
+        // Non-corner crop.
+        CropToBoundingBoxTestParams{{1, 5, 5, 1},
+                                    GetIotaVec(1, 5, 5, 1),
+                                    2,
+                                    1,
+                                    3,
+                                    3,
+                                    {1, 3, 3, 1},
+                                    {11, 12, 13,  //
+                                     16, 17, 18,  //
+                                     21, 22, 23}},
+        // Full image size crop.
+        CropToBoundingBoxTestParams{{1, 3, 3, 1},
+                                    GetIotaVec(1, 3, 3, 1),
+                                    0,
+                                    0,
+                                    3,
+                                    3,
+                                    {1, 3, 3, 1},
+                                    {0, 1, 2,  //
+                                     3, 4, 5,  //
+                                     6, 7, 8}},
+        // One-element size crop.
+        CropToBoundingBoxTestParams{{1, 3, 3, 1},
+                                    GetIotaVec(1, 3, 3, 1),
+                                    1,
+                                    1,
+                                    1,
+                                    1,
+                                    {1, 1, 1, 1},
+                                    {4}},
+        // Multichannel image.
+        CropToBoundingBoxTestParams{{1, 5, 5, 3},
+                                    GetIotaVec(1, 5, 5, 3),
+                                    2,
+                                    2,
+                                    2,
+                                    2,
+                                    {1, 2, 2, 3},
+                                    {36, 37, 38, 39, 40, 41,  //
+                                     51, 52, 53, 54, 55, 56}},
+        // Multibatch multichannel image.
+        CropToBoundingBoxTestParams{{2, 5, 5, 2},
+                                    GetIotaVec(2, 5, 5, 2),
+                                    2,
+                                    2,
+                                    2,
+                                    2,
+                                    {2, 2, 2, 2},
+                                    {24, 25, 26, 27, 34, 35, 36, 37,  //
+                                     74, 75, 76, 77, 84, 85, 86, 87}},
+    }));
+
+}  // namespace
+}  // namespace crop
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/flip_left_right.cc

@@ -0,0 +1,92 @@
+/* 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 "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace flip_left_right {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+using ::ml_adj::data::TypeWidth;
+
+// TODO(vdziuba): Move flip functionality to header.
+void FlipLeftRight(dim_t batches, dim_t input_height, dim_t input_width,
+                   const char* input_data, char* output_data,
+                   dim_t chunk_size) {
+  const ind_t row_stride = static_cast<ind_t>(input_width) * chunk_size;
+  const ind_t batch_stride = row_stride * input_height;
+
+  // Iterate over batches to flip multi-channel image.
+  for (ind_t b = 0; b < batches; ++b) {
+    const char* src_data_ptr = input_data + b * batch_stride;
+    char* dst_data_ptr = output_data + b * batch_stride;
+
+    for (ind_t y = 0; y < input_height; ++y) {
+      const char* src_ptr_row =
+          src_data_ptr + y * row_stride +
+          static_cast<ind_t>(input_width - 1) * chunk_size;
+      char* dst_ptr_row = dst_data_ptr + y * row_stride;
+      for (ind_t x = 0; x < input_width; ++x) {
+        std::memcpy(dst_ptr_row, src_ptr_row, chunk_size);
+
+        src_ptr_row -= chunk_size;
+        dst_ptr_row += chunk_size;
+      }
+    }
+  }
+}
+
+// Flips the given input horizontally. Supports any datatype.
+void ComputeFlipLeftRight(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  const char* img_data = reinterpret_cast<const char*>(img->Data());
+  const dim_t num_batches = img->Dims()[0];
+  const dim_t height = img->Dims()[1];
+  const dim_t width = img->Dims()[2];
+  const dim_t num_channels = img->Dims()[3];
+  const dim_t chunk_size = TypeWidth(img->Type()) * num_channels;
+
+  if (num_batches == 0 || height == 0 || width == 0) return;
+
+  // Resize output buffer.
+  MutableDataRef* output = outputs[0];
+  output->Resize({num_batches, height, width, num_channels});
+  char* output_data = reinterpret_cast<char*>(output->Data());
+
+  FlipLeftRight(num_batches, height, width, img_data, output_data, chunk_size);
+}
+
+}  // namespace
+
+const Algo* Impl_FlipLeftRight() {
+  static constexpr Algo kFlipLeftRight = {&ComputeFlipLeftRight, nullptr};
+  return &kFlipLeftRight;
+}
+
+}  // namespace flip_left_right
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/flip_left_right.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_LEFT_RIGHT_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_LEFT_RIGHT_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace flip_left_right {
+
+// Flip (left_right)
+//
+// Inputs: [img: any]
+// Ouputs: [img: any]
+//
+// Flips the given image horizontally (left to right).
+// Mimics semantic of `tf.image.flip_left_right`.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/flip_left_right
+
+const algo::Algo* Impl_FlipLeftRight();
+
+}  // namespace flip_left_right
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_LEFT_RIGHT_H_

tensorflow/lite/experimental/ml_adjacent/algo/flip_left_right_test.cc

@@ -0,0 +1,115 @@
+/* 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/experimental/ml_adjacent/algo/flip_left_right.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace flip_left_right {
+namespace {
+
+struct FlipLeftRightTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class FlipLeftRightTest
+    : public ::testing::TestWithParam<FlipLeftRightTestParams> {};
+
+TEST_P(FlipLeftRightTest, FloatPixelType) {
+  constexpr float kAbsError = 0.01f;
+  const FlipLeftRightTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Flip image horizontally.
+  const Algo* flip_left_right = Impl_FlipLeftRight();
+  flip_left_right->process({&img}, {&output});
+
+  // Check resize output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    FlipLeftRightTests, FlipLeftRightTest,
+    testing::ValuesIn({
+        FlipLeftRightTestParams{/*img_dims=*/{1, 3, 3, 1},
+                                /*img_data=*/
+                                {11, 12, 13,  //
+                                 21, 22, 23,  //
+                                 31, 32, 33},
+                                /*expected_data=*/
+                                {13, 12, 11,  //
+                                 23, 22, 21,  //
+                                 33, 32, 31},
+                                /*expected_shape=*/{1, 3, 3, 1}},
+        FlipLeftRightTestParams{/*img_dims=*/{1, 3, 3, 2},
+                                /*img_data=*/
+                                {11, 2, 12, 3, 13, 4,  //
+                                 21, 3, 22, 4, 23, 5,  //
+                                 31, 4, 32, 5, 33, 6},
+                                /*expected_data=*/
+                                {13, 4, 12, 3, 11, 2,  //
+                                 23, 5, 22, 4, 21, 3,  //
+                                 33, 6, 32, 5, 31, 4},
+                                /*expected_shape=*/{1, 3, 3, 2}},
+        FlipLeftRightTestParams{/*img_dims=*/{2, 3, 3, 2},
+                                /*img_data=*/
+                                {11, 2, 12, 3, 13, 4,  //
+                                 21, 3, 22, 4, 23, 5,  //
+                                 31, 4, 32, 5, 33, 6,  //
+                                 //
+                                 13, 4, 12, 3, 11, 2,  //
+                                 23, 5, 22, 4, 21, 3,  //
+                                 33, 6, 32, 5, 31, 4},
+                                /*expected_data=*/
+                                {13, 4, 12, 3, 11, 2,  //
+                                 23, 5, 22, 4, 21, 3,  //
+                                 33, 6, 32, 5, 31, 4,  //
+                                 //
+                                 11, 2, 12, 3, 13, 4,  //
+                                 21, 3, 22, 4, 23, 5,  //
+                                 31, 4, 32, 5, 33, 6},
+                                /*expected_shape=*/{2, 3, 3, 2}},
+    }));
+
+}  // namespace
+}  // namespace flip_left_right
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/flip_up_down.cc

@@ -0,0 +1,83 @@
+/* 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 "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace flip_up_down {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+using ::ml_adj::data::TypeWidth;
+
+void FlipUpDown(dim_t batches, dim_t input_height, dim_t input_width,
+                const char* input_data, char* output_data, dim_t chunk_size) {
+  const ind_t row_stride = static_cast<ind_t>(input_width) * chunk_size;
+  const ind_t batch_stride = row_stride * input_height;
+
+  // Iterate over batches to flip multi-channel image.
+  for (ind_t b = 0; b < batches; ++b) {
+    const char* src_data_ptr = input_data + b * batch_stride;
+    char* dst_data_ptr = output_data + b * batch_stride;
+
+    for (ind_t y = 0; y < input_height; ++y) {
+      const char* src_ptr_row =
+          src_data_ptr + static_cast<ind_t>(input_height - y - 1) * row_stride;
+      char* dst_ptr_row = dst_data_ptr + y * row_stride;
+      std::memcpy(dst_ptr_row, src_ptr_row, row_stride);
+    }
+  }
+}
+
+// Flips the given input vertically. Supports any datatype.
+void ComputeFlipUpDown(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  const char* img_data = reinterpret_cast<const char*>(img->Data());
+  const dim_t num_batches = img->Dims()[0];
+  const dim_t height = img->Dims()[1];
+  const dim_t width = img->Dims()[2];
+  const dim_t num_channels = img->Dims()[3];
+  const dim_t chunk_size = TypeWidth(img->Type()) * num_channels;
+
+  if (num_batches == 0 || height == 0 || width == 0) return;
+
+  // Resize output buffer.
+  MutableDataRef* output = outputs[0];
+  output->Resize({num_batches, height, width, num_channels});
+  char* output_data = reinterpret_cast<char*>(output->Data());
+
+  FlipUpDown(num_batches, height, width, img_data, output_data, chunk_size);
+}
+
+}  // namespace
+
+const Algo* Impl_FlipUpDown() {
+  static constexpr Algo kFlipUpDown = {&ComputeFlipUpDown, nullptr};
+  return &kFlipUpDown;
+}
+
+}  // namespace flip_up_down
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/flip_up_down.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_UP_DOWN_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_UP_DOWN_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace flip_up_down {
+
+// Flip (up to down)
+//
+// Inputs: [img: any]
+// Ouputs: [img: any]
+//
+// Flips the given image vertically (up to down).
+// Mimics semantic of `tf.image.flip_up_down.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/flip_up_down
+
+const algo::Algo* Impl_FlipUpDown();
+
+}  // namespace flip_up_down
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_FLIP_UP_DOWN_H_

tensorflow/lite/experimental/ml_adjacent/algo/flip_up_down_test.cc

@@ -0,0 +1,114 @@
+/* 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/experimental/ml_adjacent/algo/flip_up_down.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace flip_up_down {
+namespace {
+
+struct FlipUpDownTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class FlipUpDownTest : public ::testing::TestWithParam<FlipUpDownTestParams> {};
+
+TEST_P(FlipUpDownTest, FloatPixelType) {
+  constexpr float kAbsError = 0.01f;
+  const FlipUpDownTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Flip image vertically.
+  const Algo* flip_up_down = Impl_FlipUpDown();
+  flip_up_down->process({&img}, {&output});
+
+  // Check resize output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    FlipUpDownTests, FlipUpDownTest,
+    testing::ValuesIn({
+        FlipUpDownTestParams{/*img_dims=*/{1, 3, 3, 1},
+                             /*img_data=*/
+                             {11, 12, 13,  //
+                              21, 22, 23,  //
+                              31, 32, 33},
+                             /*expected_data=*/
+                             {31, 32, 33,  //
+                              21, 22, 23,  //
+                              11, 12, 13},
+                             /*expected_shape=*/{1, 3, 3, 1}},
+        FlipUpDownTestParams{/*img_dims=*/{1, 3, 3, 2},
+                             /*img_data=*/
+                             {11, 2, 12, 3, 13, 4,  //
+                              21, 3, 22, 4, 23, 5,  //
+                              31, 4, 32, 5, 33, 6},
+                             /*expected_data=*/
+                             {31, 4, 32, 5, 33, 6,  //
+                              21, 3, 22, 4, 23, 5,  //
+                              11, 2, 12, 3, 13, 4},
+                             /*expected_shape=*/{1, 3, 3, 2}},
+        FlipUpDownTestParams{/*img_dims=*/{2, 3, 3, 2},
+                             /*img_data=*/
+                             {11, 2, 12, 3, 13, 4,  //
+                              21, 3, 22, 4, 23, 5,  //
+                              31, 4, 32, 5, 33, 6,  //
+                              //
+                              13, 4, 12, 3, 11, 2,  //
+                              23, 5, 22, 4, 21, 3,  //
+                              33, 6, 32, 5, 31, 4},
+                             /*expected_data=*/
+                             {31, 4, 32, 5, 33, 6,  //
+                              21, 3, 22, 4, 23, 5,  //
+                              11, 2, 12, 3, 13, 4,  //
+                              //
+                              33, 6, 32, 5, 31, 4,  //
+                              23, 5, 22, 4, 21, 3,  //
+                              13, 4, 12, 3, 11, 2},
+                             /*expected_shape=*/{2, 3, 3, 2}},
+    }));
+
+}  // namespace
+}  // namespace flip_up_down
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/image_utils.h

@@ -0,0 +1,53 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_IMAGE_UTILS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_IMAGE_UTILS_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+
+inline void ConvertColorSpace(dim_t batches, dim_t height, dim_t width,
+                              const float* input_data, float* output_data,
+                              const float* kernel_data, dim_t kernel_size) {
+  TFLITE_CHECK_EQ(kernel_size, 9);
+
+  const dim_t num_channels = 3;
+  const ind_t output_num_pixels = static_cast<ind_t>(batches) * width * height;
+  const float* src_data_ptr = input_data;
+  float* dst_data_ptr = output_data;
+
+  for (ind_t i = 0; i < output_num_pixels; ++i) {
+    dst_data_ptr[0] = kernel_data[0] * src_data_ptr[0] +
+                      kernel_data[1] * src_data_ptr[1] +
+                      kernel_data[2] * src_data_ptr[2];
+
+    dst_data_ptr[1] = kernel_data[3] * src_data_ptr[0] +
+                      kernel_data[4] * src_data_ptr[1] +
+                      kernel_data[5] * src_data_ptr[2];
+
+    dst_data_ptr[2] = kernel_data[6] * src_data_ptr[0] +
+                      kernel_data[7] * src_data_ptr[1] +
+                      kernel_data[8] * src_data_ptr[2];
+
+    src_data_ptr += num_channels;
+    dst_data_ptr += num_channels;
+  }
+}
+
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_IMAGE_UTILS_H_

tensorflow/lite/experimental/ml_adjacent/algo/per_image_standardization.cc

@@ -0,0 +1,116 @@
+/* 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/experimental/ml_adjacent/algo/per_image_standardization.h"
+
+#include <cmath>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace per_image_standardization {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+inline void PerImageStandardization(dim_t batches, dim_t height, dim_t width,
+                                    dim_t num_channels, const float* input_data,
+                                    float* output_data) {
+  const ind_t num_pixels_per_image =
+      static_cast<ind_t>(height) * width * num_channels;
+
+  if (batches == 0 || num_pixels_per_image == 0) return;
+
+  const float inv_num_pixels_per_image = 1.0f / num_pixels_per_image;
+
+  for (ind_t b = 0; b < batches; ++b) {
+    const ind_t offset = b * num_pixels_per_image;
+    const float* input_ptr = input_data + offset;
+    float* output_ptr = output_data + offset;
+
+    // Compute mean of all image pixels.
+    float mean = 0.0f;
+    for (ind_t i = 0; i < num_pixels_per_image; ++i) {
+      mean += input_ptr[i];
+    }
+    mean *= inv_num_pixels_per_image;
+
+    // Compute adjusted_stddev = max(stddev, 1.0/sqrt(N)) which is capped away
+    // from 0 to protect against division by 0 when handling uniform images.
+    // Reusing computed (x - mean) and saving it to the output here for speed
+    // optimization.
+    float variance = 0.0f;
+    for (ind_t i = 0; i < num_pixels_per_image; ++i) {
+      const float diff = input_ptr[i] - mean;
+      variance += diff * diff * inv_num_pixels_per_image;
+      output_ptr[i] = diff;
+    }
+
+    // Divide output by adjusted stddev.
+    const float adjusted_stddev =
+        std::fmax(std::sqrt(variance),
+                  1.0f / std::sqrt(static_cast<float>(num_pixels_per_image)));
+    const float inv_adjusted_stddev = 1.0f / adjusted_stddev;
+    for (ind_t i = 0; i < num_pixels_per_image; ++i) {
+      output_ptr[i] *= inv_adjusted_stddev;
+    }
+  }
+}
+
+// Linearly scales each image in input to have mean 0 and variance 1. Works for
+// float datatype.
+void ComputePerImageStandardization(const InputPack& inputs,
+                                    const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  MutableDataRef* output = outputs[0];
+
+  TFLITE_CHECK_EQ(img->Type(), etype_t::f32);
+  TFLITE_CHECK_EQ(output->Type(), etype_t::f32);
+
+  const float* img_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t img_num_batches = img->Dims()[0];
+  const dim_t img_height = img->Dims()[1];
+  const dim_t img_width = img->Dims()[2];
+  const dim_t img_num_channels = img->Dims()[3];
+
+  if (img_num_batches == 0 || img_height == 0 || img_width == 0) return;
+
+  // Resize output buffer for resized image.
+  output->Resize({img_num_batches, img_height, img_width, img_num_channels});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  PerImageStandardization(img_num_batches, img_height, img_width,
+                          img_num_channels, img_data, output_data);
+}
+
+}  // namespace
+
+const Algo* Impl_PerImageStandardization() {
+  static constexpr Algo kPerImageStandardization = {
+      &ComputePerImageStandardization, nullptr};
+  return &kPerImageStandardization;
+}
+
+}  // namespace per_image_standardization
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/per_image_standardization.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_PER_IMAGE_STANDARDIZATION_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_PER_IMAGE_STANDARDIZATION_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace per_image_standardization {
+
+// Per Image Standardization
+//
+// Inputs: [img: float]
+// Ouputs: [img: float]
+//
+// Linearly scales each image in input to have mean 0 and variance 1.
+// Mimics semantic of `tf.image.per_image_standardization`.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/per_image_standardization
+
+const algo::Algo* Impl_PerImageStandardization();
+
+}  // namespace per_image_standardization
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_PER_IMAGE_STANDARDIZATION_H_

tensorflow/lite/experimental/ml_adjacent/algo/per_image_standardization_test.cc

@@ -0,0 +1,141 @@
+/* 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/experimental/ml_adjacent/algo/per_image_standardization.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace per_image_standardization {
+namespace {
+
+struct PerImageStandardizationTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+};
+
+class PerImageStandardizationTest
+    : public testing::TestWithParam<PerImageStandardizationTestParams> {};
+
+TEST_P(PerImageStandardizationTest, FloatPixelType) {
+  const PerImageStandardizationTestParams& params = GetParam();
+
+  // Create input image tensor.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Create empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Run per image standardization custom call.
+  const Algo* per_image_standardization = Impl_PerImageStandardization();
+  per_image_standardization->process({&img}, {&output});
+
+  // Check output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), img.Dims());
+  constexpr float kAbsError = 0.01f;
+  float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    // TODO: b/298483909 - This print is being used in multiple locations.
+    // Ideally should be moved to a shared test_utils file.
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    PerImageStandardizationTests, PerImageStandardizationTest,
+    testing::ValuesIn({
+        // 4-D Tensor of shape [batch, height, width, channels] is used below.
+        PerImageStandardizationTestParams{/*img_dims=*/
+                                          {1, 2, 2, 1},
+                                          /*img_data=*/
+                                          {1, 2,  //
+                                           3, 4},
+                                          /*expected_data=*/
+                                          {-1.3416407, -0.4472136,  //
+                                           0.4472136, 1.3416407}},
+        // Two images in batch.
+        PerImageStandardizationTestParams{/*img_dims=*/
+                                          {2, 2, 2, 1},
+                                          /*img_data=*/
+                                          {1, 2,  //
+                                           3, 4,  //
+                                                  //
+                                           1, 2,  //
+                                           4, 8},
+                                          /*expected_data=*/
+                                          {-1.3416407, -0.4472136,   //
+                                           0.4472136, 1.3416407,     //
+                                                                     //
+                                           -1.0257553, -0.65275335,  //
+                                           0.09325048, 1.5852581}},
+        // Multi-channel multi-batch image.
+        PerImageStandardizationTestParams{/*img_dims=*/
+                                          {2, 2, 2, 2},
+                                          /*img_data=*/
+                                          {1, 2,  //
+                                           1, 3,  //
+                                           1, 4,  //
+                                           1, 5,  //
+                                                  //
+                                           1, 2,  //
+                                           2, 2,  //
+                                           3, 2,  //
+                                           4, 2},
+                                          /*expected_data=*/
+                                          {-0.8451542, -0.16903085,   //
+                                           -0.8451542, 0.50709254,    //
+                                           -0.8451542, 1.1832159,     //
+                                           -0.8451542, 1.8593392,     //
+                                                                      //
+                                           -1.5075567, -0.30151135,   //
+                                           -0.30151135, -0.30151135,  //
+                                           0.904534, -0.30151135,     //
+                                           2.1105793, -0.30151135}},
+    }));
+
+TEST(PerImageStandardizationDeathTest, WrongInputTypeFails) {
+  OwningVectorRef img(etype_t::i32);
+  img.Resize({1, 1, 1, 1});
+  OwningVectorRef output(etype_t::f32);
+
+  const Algo* per_image_standardization = Impl_PerImageStandardization();
+  EXPECT_DEATH(per_image_standardization->process({&img}, {&output}), "");
+}
+
+TEST(PerImageStandardizationDeathTest, WrongOutputTypeFails) {
+  OwningVectorRef img(etype_t::f32);
+  img.Resize({1, 1, 1, 1});
+  OwningVectorRef output(etype_t::i32);
+
+  const Algo* per_image_standardization = Impl_PerImageStandardization();
+  EXPECT_DEATH(per_image_standardization->process({&img}, {&output}), "");
+}
+}  // namespace
+}  // namespace per_image_standardization
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/resize.cc

@@ -0,0 +1,225 @@
+/* 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 <cmath>
+#include <cstdint>
+#include <cstring>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace resize {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+// Calculates bilinear interpolated values and lower and upper bounds.
+inline void ComputeInterpolationValues(const float value, const float scale,
+                                       int32_t input_size, float& scaled_value,
+                                       int32_t& lower_bound,
+                                       int32_t& upper_bound) {
+  scaled_value = value * scale;
+
+  float scaled_value_floor = std::floor(scaled_value);
+  lower_bound = std::max(static_cast<int32_t>(scaled_value_floor), 0);
+  upper_bound =
+      std::min(static_cast<int32_t>(std::ceil(scaled_value)), input_size - 1);
+}
+
+// Applies depth-wise scaling.
+inline void ScaleDepthwise(const float* input_ptr, int32_t depth, float scale,
+                           float* output_ptr) {
+  for (int32_t i = 0; i < depth; ++i) {
+    *output_ptr += *input_ptr * scale;
+    output_ptr++;
+    input_ptr++;
+  }
+}
+
+// Calculates 1D pixel offset in NHWC data.
+inline int64_t Offset(const int32_t* dims_data, int32_t dims_num, int64_t i0,
+                      int64_t i1, int64_t i2, int64_t i3) {
+  TFLITE_CHECK_EQ(dims_num, 3);
+  return ((i0 * dims_data[0] + i1) * dims_data[1] + i2) * dims_data[2] + i3;
+}
+
+// Generic implementation of bilinear resize.
+inline void ResizeBilinear(int32_t batches, int32_t input_height,
+                           int32_t input_width, int32_t depth,
+                           int32_t output_height, int32_t output_width,
+                           float height_scale, float width_scale,
+                           const float* input_data, float* output_data) {
+  memset(output_data, 0,
+         static_cast<size_t>(batches) * output_height * output_width * depth *
+             sizeof(float));
+  const int dims_data[] = {input_height, input_width, depth};
+  const int dims_num = sizeof(dims_data) / sizeof(dims_data[0]);
+
+  int64_t output_offset = 0;
+  for (int b = 0; b < batches; ++b) {
+    for (int y = 0; y < output_height; ++y) {
+      float input_y = 0.0f;
+      int32_t y0 = 0;
+      int32_t y1 = 0;
+      ComputeInterpolationValues(y, height_scale, input_height, input_y, y0,
+                                 y1);
+      for (int x = 0; x < output_width; ++x) {
+        float input_x = 0.0f;
+        int32_t x0 = 0;
+        int32_t x1 = 0;
+        ComputeInterpolationValues(x, width_scale, input_width, input_x, x0,
+                                   x1);
+        float* output_ptr = output_data + output_offset;
+
+        // Run kernel on the 4 corners of the bilinear resize algorithm.
+        int64_t input_offset = Offset(dims_data, dims_num, b, y0, x0, 0);
+        float scale = (1 - (input_y - y0)) * (1 - (input_x - x0));
+        const float* input_ptr = &input_data[input_offset];
+        ScaleDepthwise(input_ptr, depth, scale, output_ptr);
+
+        input_offset = Offset(dims_data, dims_num, b, y0, x1, 0);
+        scale = (1 - (input_y - y0)) * (input_x - x0);
+        input_ptr = input_data + input_offset;
+        ScaleDepthwise(input_ptr, depth, scale, output_ptr);
+
+        input_offset = Offset(dims_data, dims_num, b, y1, x0, 0);
+        scale = (input_y - y0) * (1 - (input_x - x0));
+        input_ptr = &input_data[input_offset];
+        ScaleDepthwise(input_ptr, depth, scale, output_ptr);
+
+        input_offset = Offset(dims_data, dims_num, b, y1, x1, 0);
+        scale = (input_y - y0) * (input_x - x0);
+        input_ptr = &input_data[input_offset];
+        ScaleDepthwise(input_ptr, depth, scale, output_ptr);
+
+        output_offset += depth;
+      }
+    }
+  }
+}
+
+// Optimized implementation of bilinear resize for 2X upscaling case.
+inline void ResizeBilinear2x2(int32_t batches, int32_t input_height,
+                              int32_t input_width, int32_t depth,
+                              int32_t output_height, int32_t output_width,
+                              const float* input_data, float* output_data) {
+  const int input_dims_data[] = {input_height, input_width, depth};
+  const int output_dims_data[] = {output_height, output_width, depth};
+  const int dims_num = sizeof(input_dims_data) / sizeof(input_dims_data[0]);
+
+  for (int32_t b = 0; b < batches; ++b) {
+    for (int32_t y0 = 0, y = 0; y <= output_height - 2; y += 2, ++y0) {
+      for (int32_t x0 = 0, x = 0; x <= output_width - 2; x += 2, ++x0) {
+        int32_t x1 = std::min(x0 + 1, input_width - 1);
+        int32_t y1 = std::min(y0 + 1, input_height - 1);
+
+        const int64_t input_x_offset = static_cast<int64_t>(x1 - x0) * depth;
+        const int64_t input_y_offset =
+            static_cast<int64_t>(y1 - y0) * depth * input_width;
+        const int64_t output_x_offset = depth;
+        const int64_t output_y_offset =
+            static_cast<int64_t>(depth) * output_width;
+
+        for (int ch = 0; ch < depth; ++ch) {
+          const int64_t input_offset =
+              Offset(input_dims_data, dims_num, b, y0, x0, ch);
+
+          const float x0y0 = input_data[input_offset];
+          const float x1y0 = input_data[input_offset + input_x_offset];
+          const float x0y1 = input_data[input_offset + input_y_offset];
+          const float x1y1 =
+              input_data[input_offset + input_x_offset + input_y_offset];
+
+          // Calculate top left corner value.
+          const int64_t output_offset =
+              Offset(output_dims_data, dims_num, b, y, x, ch);
+          output_data[output_offset] = x0y0;
+
+          // Calculate top right corner value.
+          output_data[output_offset + output_x_offset] = (x0y0 + x1y0) * 0.5f;
+
+          // Calculate bottom left corner value.
+          float output = (x0y0 + x0y1) * 0.5f;
+          output_data[output_offset + output_y_offset] = output;
+
+          // Calculate bottom right corner value.
+          output_data[output_offset + output_x_offset + output_y_offset] =
+              (output + ((x1y0 + x1y1) * 0.5f)) * 0.5f;
+        }
+      }
+    }
+  }
+}
+
+// Resizes given input. Supports `float` datatype only.
+void ComputeResize(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 2);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  const float* img_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t img_num_batches = img->Dims()[0];
+  const dim_t img_height = img->Dims()[1];
+  const dim_t img_width = img->Dims()[2];
+  const dim_t img_num_channels = img->Dims()[3];
+
+  // Extract new image size.
+  const DataRef* size = inputs[1];
+  TFLITE_CHECK_EQ(size->NumElements(), 2);
+  const dim_t* size_data = reinterpret_cast<const dim_t*>(size->Data());
+  const dim_t new_height = size_data[0];
+  const dim_t new_width = size_data[1];
+
+  if (img_num_batches == 0 || img_num_channels == 0 || new_height == 0 ||
+      new_width == 0) {
+    return;
+  }
+  TFLITE_CHECK(img_height > 0 && img_width > 0);
+
+  // Resize output buffer for resized image.
+  MutableDataRef* output = outputs[0];
+  output->Resize({img_num_batches, new_height, new_width, img_num_channels});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  const float width_scale = static_cast<float>(img_width) / new_width;
+  const float height_scale = static_cast<float>(img_height) / new_height;
+
+  if (new_width == 2 * img_width && new_height == 2 * img_height) {
+    ResizeBilinear2x2(img_num_batches, img_height, img_width, img_num_channels,
+                      new_height, new_width, img_data, output_data);
+    return;
+  }
+
+  ResizeBilinear(img_num_batches, img_height, img_width, img_num_channels,
+                 new_height, new_width, height_scale, width_scale, img_data,
+                 output_data);
+}
+
+}  // namespace
+
+const Algo* Impl_Resize() {
+  static constexpr Algo kResize = {&ComputeResize, nullptr};
+  return &kResize;
+}
+
+}  // namespace resize
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/resize.h

@@ -0,0 +1,39 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_RESIZE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RESIZE_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace resize {
+
+// Bilinear Resize
+//
+// Inputs: [img: float, new size: vector<unsigned>]
+// Ouputs: [img: float]
+//
+// Resizes the given image, scaling height and width dimensions to `new size`.
+// Mimics semantic of `tf.image.resize`, where `method` equals to bilinear
+// interpolation.
+//
+// https://www.tensorflow.org/api_docs/python/tf/image/resize
+
+const algo::Algo* Impl_Resize();
+
+}  // namespace resize
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RESIZE_H_

tensorflow/lite/experimental/ml_adjacent/algo/resize_test.cc

@@ -0,0 +1,158 @@
+/* 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/experimental/ml_adjacent/algo/resize.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace resize {
+namespace {
+
+struct ResizeTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<dim_t> size;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class ResizeTest : public ::testing::TestWithParam<ResizeTestParams> {};
+
+TEST_P(ResizeTest, FloatPixelType) {
+  constexpr float kAbsError = 0.01f;
+  const ResizeTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // New size input.
+  OwningVectorRef size(etype_t::i32);
+  size.Resize({2});
+  ASSERT_EQ(size.Bytes(), params.size.size() * sizeof(int));
+  std::memcpy(size.Data(), params.size.data(), size.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Run image resize custom call.
+  const Algo* resize = Impl_Resize();
+  resize->process({&img, &size}, {&output});
+
+  // Check resize output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ResizeTests, ResizeTest,
+    testing::ValuesIn({
+        // 4-D Tensor of shape [batch, height, width, channels] is used below.
+
+        // Check 2x2 image resizing to 4x4 size.
+        ResizeTestParams{/*img_dims=*/{1, 2, 2, 1},
+                         /*img_data=*/
+                         {1, 2,  //
+                          3, 4},
+                         /*new_shape=*/{4, 4},
+                         /*expected_data=*/
+                         {1, 1.5, 2, 2,  //
+                          2, 2.5, 3, 3,  //
+                          3, 3.5, 4, 4,  //
+                          3, 3.5, 4, 4},
+                         /*expected_shape=*/{1, 4, 4, 1}},
+
+        // Check vertical resizing of 2-element vector.
+        ResizeTestParams{/*img_dims=*/{1, 2, 1, 1},
+                         /*img_data=*/{3, 9},
+                         /*size=*/{3, 1},
+                         /*expected_data=*/{3, 7, 9},
+                         /*expected_shape=*/{1, 3, 1, 1}},
+
+        // Check horizontal resizing of 2-element vector.
+        ResizeTestParams{/*img_dims=*/{1, 1, 2, 1},
+                         /*img_data=*/{3, 6},
+                         /*size=*/{1, 3},
+                         /*expected_data=*/{3, 5, 6},
+                         /*expected_shape=*/{1, 1, 3, 1}},
+
+        // Check 2x2 image resizing.
+        ResizeTestParams{/*img_dims=*/{1, 2, 2, 1},
+                         /*img_data=*/
+                         {3, 6,  //
+                          9, 12},
+                         /*size=*/{3, 3},
+                         /*expected_data=*/
+                         {3, 5, 6,   //
+                          7, 9, 10,  //
+                          9, 11, 12},
+                         /*expected_shape=*/{1, 3, 3, 1}},
+
+        // Check 2x2 image resizing (2 images in the batch).
+        ResizeTestParams{/*img_dims=*/{2, 2, 2, 1},
+                         /*img_data=*/
+                         {3, 6,   //
+                          9, 12,  //
+                                  //
+                          4, 10,  //
+                          10, 16},
+                         /*size=*/{3, 3},
+                         /*expected_data=*/
+                         {3, 5, 6,    //
+                          7, 9, 10,   //
+                          9, 11, 12,  //
+                                      //
+                          4, 8, 10,   //
+                          8, 12, 14,  //
+                          10, 14, 16},
+                         /*expected_shape=*/{2, 3, 3, 1}},
+
+        // Check 2x2 with 2 channels image resizing.
+        ResizeTestParams{
+            /*img_dims=*/{1, 2, 2, 2},
+            /*img_data=*/{3, 4, 6, 10, 9, 10, 12, 16},
+            /*size=*/{3, 3},
+            /*expected_data=*/
+            {3, 4, 5, 8, 6, 10, 7, 8, 9, 12, 10, 14, 9, 10, 11, 14, 12, 16},
+            /*expected_shape=*/{1, 3, 3, 2}},
+
+        // Check horizontal resizing of 2-element vector with large values.
+        ResizeTestParams{/*img_dims=*/{1, 1, 2, 1},
+                         /*img_data=*/{32765, 32767},
+                         /*size=*/{1, 3},
+                         /*expected_data=*/{32765, 32766.33f, 32767},
+                         /*expected_shape=*/{1, 1, 3, 1}},
+    }));
+
+}  // namespace
+}  // namespace resize
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_grayscale.cc

@@ -0,0 +1,85 @@
+/* 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/experimental/ml_adjacent/algo/rgb_to_grayscale.h"
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace rgb_to_grayscale {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+inline void ConvertRgbToGrayscale(dim_t batches, dim_t height, dim_t width,
+                                  const float* input_data, float* output_data) {
+  const ind_t output_num_pixels = static_cast<ind_t>(batches) * width * height;
+  // Reference for converting between RGB and grayscale. Same as in
+  // tf.image.rgb_to_grayscale: https://en.wikipedia.org/wiki/Luma_%28video%29.
+  static constexpr float kRgb2GrayscaleKernel[] = {0.2989f, 0.5870f, 0.1140f};
+  const float* src_ptr = input_data;
+  float* dst_ptr = output_data;
+  for (ind_t i = 0; i < output_num_pixels; ++i) {
+    *dst_ptr = kRgb2GrayscaleKernel[0] * src_ptr[0] +
+               kRgb2GrayscaleKernel[1] * src_ptr[1] +
+               kRgb2GrayscaleKernel[2] * src_ptr[2];
+    src_ptr += 3;  // Step is number of input channels, which is 3 for RGB.
+    dst_ptr++;
+  }
+}
+
+// Converts each image in input from RGB to grayscale. Works for float datatype.
+void ComputeRgbToGrayscale(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  MutableDataRef* output = outputs[0];
+
+  TFLITE_CHECK_EQ(img->Type(), etype_t::f32);
+  TFLITE_CHECK_EQ(output->Type(), etype_t::f32);
+
+  const float* img_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t img_num_batches = img->Dims()[0];
+  const dim_t img_height = img->Dims()[1];
+  const dim_t img_width = img->Dims()[2];
+  const dim_t channels = img->Dims()[3];
+  TFLITE_CHECK_EQ(channels, 3);
+
+  if (img_num_batches == 0 || img_height == 0 || img_width == 0) return;
+
+  // Resize output buffer for single-channel output image.
+  output->Resize({img_num_batches, img_height, img_width, 1});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  ConvertRgbToGrayscale(img_num_batches, img_height, img_width, img_data,
+                        output_data);
+}
+
+}  // namespace
+
+const Algo* Impl_RgbToGrayscale() {
+  static constexpr Algo kRgbToGrayscale = {&ComputeRgbToGrayscale, nullptr};
+  return &kRgbToGrayscale;
+}
+
+}  // namespace rgb_to_grayscale
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_grayscale.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_GRAYSCALE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_GRAYSCALE_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace rgb_to_grayscale {
+
+// RGB to Grayscale conversion
+//
+// Inputs: [img: float]
+// Ouputs: [img: float]
+//
+// Converts one or more images from RGB to Grayscale.
+// Mimics semantic of `tf.image.rgb_to_grayscale`.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/rgb_to_grayscale
+
+const algo::Algo* Impl_RgbToGrayscale();
+
+}  // namespace rgb_to_grayscale
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_GRAYSCALE_H_

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_grayscale_test.cc

@@ -0,0 +1,119 @@
+/* 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/experimental/ml_adjacent/algo/rgb_to_grayscale.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace rgb_to_grayscale {
+namespace {
+
+struct RgbToGrayscaleTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class RgbToGrayscaleTest
+    : public ::testing::TestWithParam<RgbToGrayscaleTestParams> {};
+
+TEST_P(RgbToGrayscaleTest, FloatPixelType) {
+  const RgbToGrayscaleTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Empty output.
+  OwningVectorRef output(etype_t::f32);
+
+  // Convert RGB to grayscale.
+  const Algo* rgb_to_grayscale = Impl_RgbToGrayscale();
+  rgb_to_grayscale->process({&img}, {&output});
+
+  // Check output size and values.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  constexpr float kAbsError = 0.1f;
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    RgbToGrayscaleTests, RgbToGrayscaleTest,
+    testing::ValuesIn({
+        RgbToGrayscaleTestParams{/*img_dims=*/{1, 3, 2, 3},
+                                 /*img_data=*/
+                                 {11, 111, 211,  //
+                                  12, 112, 212,  //
+                                  //
+                                  21, 121, 221,  //
+                                  22, 122, 222,  //
+                                  //
+                                  31, 131, 231,  //
+                                  32, 132, 232},
+                                 /*expected_data=*/
+                                 {92.5f, 93.5f, 102.5f,  //
+                                  103.5f, 112.5f, 113.5f},
+                                 /*expected_shape=*/{1, 3, 2, 1}},
+        // Two images in batch.
+        RgbToGrayscaleTestParams{/*img_dims=*/{2, 3, 2, 3},
+                                 /*img_data=*/
+                                 {11, 111, 211,  //
+                                  12, 112, 212,  //
+                                  //
+                                  21, 121, 221,  //
+                                  22, 122, 222,  //
+                                  //
+                                  31, 131, 231,  //
+                                  32, 132, 232,  //
+                                  //
+                                  //
+                                  51, 311, 411,  //
+                                  52, 312, 412,  //
+                                  //
+                                  61, 321, 421,  //
+                                  62, 322, 422,  //
+                                  //
+                                  71, 331, 431,  //
+                                  72, 332, 432},
+                                 /*expected_data=*/
+                                 {92.5f, 93.5f, 102.5f,    //
+                                  103.5f, 112.5f, 113.5f,  //
+                                  //
+                                  244.7f, 245.7f, 254.7f,  //
+                                  255.7f, 264.7f, 265.7f},
+                                 /*expected_shape=*/{2, 3, 2, 1}},
+    }));
+
+}  // namespace
+}  // namespace rgb_to_grayscale
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_yuv.cc

@@ -0,0 +1,73 @@
+/* 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/experimental/ml_adjacent/algo/image_utils.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace rgb_to_yuv {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+constexpr float kRgb2YuvKernel[] = {0.299f,       0.587f,       0.114f,       //
+                                    -0.14714119f, -0.28886916f, 0.43601035f,  //
+                                    0.61497538f,  -0.51496512f, -0.10001026f};
+constexpr int kRgb2YuvKernelSize =
+    sizeof(kRgb2YuvKernel) / sizeof(kRgb2YuvKernel[0]);
+
+// Converts RGB to YUV. Supports `float` datatype only.
+void ComputeRgbToYuv(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  MutableDataRef* output = outputs[0];
+
+  TFLITE_CHECK_EQ(img->Type(), etype_t::f32);
+  TFLITE_CHECK_EQ(output->Type(), etype_t::f32);
+
+  const float* input_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t batches = img->Dims()[0];
+  const dim_t height = img->Dims()[1];
+  const dim_t width = img->Dims()[2];
+  const dim_t channels = img->Dims()[3];
+  TFLITE_CHECK_EQ(channels, 3);
+
+  if (batches == 0 || height == 0 || width == 0) return;
+
+  // Resize output buffer.
+  output->Resize({batches, height, width, channels});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  ConvertColorSpace(batches, height, width, input_data, output_data,
+                    &kRgb2YuvKernel[0], kRgb2YuvKernelSize);
+}
+
+}  // namespace
+
+const Algo* Impl_RgbToYuv() {
+  static constexpr Algo kRgbToYuv = {&ComputeRgbToYuv, nullptr};
+  return &kRgbToYuv;
+}
+
+}  // namespace rgb_to_yuv
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_yuv.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_YUV_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_YUV_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace rgb_to_yuv {
+
+// RGB to YUV conversion
+//
+// Inputs: [img: float]
+// Ouputs: [img: float]
+//
+// Coverts the given 3-channel RGB image to 3-channel YUV image.
+// Mimics semantic of `tf.image.rgb_to_yuv.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/rgb_to_yuv
+
+const algo::Algo* Impl_RgbToYuv();
+
+}  // namespace rgb_to_yuv
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_RGB_TO_YUV_H_

tensorflow/lite/experimental/ml_adjacent/algo/rgb_to_yuv_test.cc

@@ -0,0 +1,119 @@
+/* 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/experimental/ml_adjacent/algo/rgb_to_yuv.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace rgb_to_yuv {
+namespace {
+
+struct Rgb2YuvTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class Rgb2YuvTest : public ::testing::TestWithParam<Rgb2YuvTestParams> {};
+
+TEST_P(Rgb2YuvTest, FloatPixelType) {
+  constexpr float kAbsError = 0.1f;
+  const Rgb2YuvTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Convert RGB to YUV image.
+  const Algo* rgb_to_yuv = Impl_RgbToYuv();
+  rgb_to_yuv->process({&img}, {&output});
+
+  // Check resize output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Rgb2YuvTests, Rgb2YuvTest,
+    testing::ValuesIn({
+        Rgb2YuvTestParams{/*img_dims=*/{1, 3, 2, 3},
+                          /*img_data=*/
+                          {11, 111, 211, 12, 112, 212,  //
+                           21, 121, 221, 22, 122, 222,  //
+                           31, 131, 231, 32, 132, 232},
+                          /*expected_data=*/
+                          {
+                              92.5f,
+                              58.3f,
+                              -71.5f,
+                              93.5f,
+                              58.3f,
+                              -71.5f,  //
+                              102.5f,
+                              58.3f,
+                              -71.5f,
+                              103.5f,
+                              58.3f,
+                              -71.5f,  //
+                              112.5f,
+                              58.3f,
+                              -71.5f,
+                              113.5f,
+                              58.3f,
+                              -71.5f,
+                          },
+                          /*expected_shape=*/{1, 3, 2, 3}},
+        Rgb2YuvTestParams{/*img_dims=*/{2, 3, 2, 3},
+                          /*img_data=*/
+                          {11, 111, 211, 12, 112, 212,  //
+                           21, 121, 221, 22, 122, 222,  //
+                           31, 131, 231, 32, 132, 232,  //
+                           11, 111, 211, 12, 112, 212,  //
+                           21, 121, 221, 22, 122, 222,  //
+                           31, 131, 231, 32, 132, 232},
+                          /*expected_data=*/
+                          {92.5f,  58.3f, -71.5f, 93.5f,  58.3f, -71.5f,  //
+                           102.5f, 58.3f, -71.5f, 103.5f, 58.3f, -71.5f,  //
+                           112.5f, 58.3f, -71.5f, 113.5f, 58.3f, -71.5f,  //
+                           92.5f,  58.3f, -71.5f, 93.5f,  58.3f, -71.5f,  //
+                           102.5f, 58.3f, -71.5f, 103.5f, 58.3f, -71.5f,  //
+                           112.5f, 58.3f, -71.5f, 113.5f, 58.3f, -71.5f},
+                          /*expected_shape=*/{2, 3, 2, 3}},
+    }));
+
+}  // namespace
+}  // namespace rgb_to_yuv
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rotate.cc

@@ -0,0 +1,336 @@
+/* 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 <cmath>
+#include <cstdint>
+#include <cstring>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace rotate {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+inline float DegreesToRadians(int angle) { return angle * M_PI / 180; }
+
+void ComputeNewSize(dim_t src_width, dim_t src_height, int angle,
+                    dim_t& dst_width, dim_t& dst_height) {
+  // Keep same size for 180 degree case.
+  dst_width = src_width;
+  dst_height = src_height;
+
+  if (angle % 90 == 0) {
+    // Define new size for 90 or 270 degree angle.
+    if (angle == 90 || angle == 270) {
+      dst_width = src_height;
+      dst_height = src_width;
+    }
+  } else {
+    // Calculate new size for arbitrary angle.
+    const float angle_rad = DegreesToRadians(angle);
+    const float cos_angle = std::cos(angle_rad);
+    const float sin_angle = std::sin(angle_rad);
+
+    const int64_t edge_x = static_cast<int64_t>(src_width) / 2;
+    const int64_t edge_y = static_cast<int64_t>(src_height) / 2;
+    for (int64_t y : {-edge_y, edge_y}) {
+      for (int64_t x : {-edge_x, edge_x}) {
+        const int64_t x_transformed =
+            static_cast<int64_t>(std::floor(cos_angle * x + sin_angle * y));
+        const int64_t y_transformed =
+            static_cast<int64_t>(std::floor(-sin_angle * x + cos_angle * y));
+
+        if (std::abs(x_transformed) > static_cast<int64_t>(dst_width) / 2) {
+          dst_width = 2 * std::abs(x_transformed);
+        }
+        if (std::abs(y_transformed) > static_cast<int64_t>(dst_height) / 2) {
+          dst_height = 2 * std::abs(y_transformed);
+        }
+      }
+    }
+  }
+}
+
+// Rotates image for 90 degree.
+void Rotate90(dim_t batches, dim_t input_height, dim_t input_width, dim_t depth,
+              dim_t output_height, dim_t output_width, const float* input_data,
+              float* output_data) {
+  TFLITE_CHECK(input_data != nullptr);
+  TFLITE_CHECK(output_data != nullptr);
+
+  const int64_t pixel_stride = depth;
+  const int64_t src_row_stride = static_cast<int64_t>(input_width) * depth;
+  const int64_t dst_row_stride = static_cast<int64_t>(output_width) * depth;
+  const int64_t src_batch_stride = src_row_stride * input_height;
+  const int64_t dst_batch_stride = dst_row_stride * output_height;
+
+  // Iterate over batches to perform the following transformation:
+  // dst[x][width - y - 1] = src[y][x].
+  for (dim_t b = 0; b < batches; ++b) {
+    const float* src_data_ptr = input_data + b * src_batch_stride;
+    float* dst_data_ptr = output_data + b * dst_batch_stride;
+
+    for (dim_t y = 0; y < input_height; ++y) {
+      const float* src_ptr_row = src_data_ptr + y * src_row_stride;
+      for (dim_t x = 0; x < input_width; ++x) {
+        float* dst_ptr_row = dst_data_ptr + x * dst_row_stride;
+
+        const float* src_ptr_pixel = src_ptr_row + x * pixel_stride;
+        float* dst_pixel_ptr =
+            dst_ptr_row + (output_width - y - 1) * pixel_stride;
+
+        for (dim_t c = 0; c < depth; ++c) {
+          *dst_pixel_ptr++ = *src_ptr_pixel++;
+        }
+      }
+    }
+  }
+}
+
+// Rotates image for 180 degree.
+void Rotate180(dim_t batches, dim_t input_height, dim_t input_width,
+               dim_t depth, dim_t output_height, dim_t output_width,
+               const float* input_data, float* output_data) {
+  TFLITE_CHECK(input_data != nullptr);
+  TFLITE_CHECK(output_data != nullptr);
+
+  const int64_t dst_pixel_stride = depth;
+  const int64_t src_row_stride = static_cast<int64_t>(input_width) * depth;
+  const int64_t dst_row_stride = static_cast<int64_t>(output_width) * depth;
+  const int64_t src_batch_stride = src_row_stride * input_height;
+  const int64_t dst_batch_stride = dst_row_stride * output_height;
+
+  // Iterate over batches to perform the following transformation:
+  // dst[height - y - 1][width - x - 1] = src[y][x].
+  for (dim_t b = 0; b < batches; ++b) {
+    const float* src_data_ptr = input_data + b * src_batch_stride;
+    float* dst_data_ptr = output_data + b * dst_batch_stride;
+
+    for (dim_t y = 0; y < input_height; ++y) {
+      const float* src_ptr_row = src_data_ptr + y * src_row_stride;
+      float* dst_ptr_row = dst_data_ptr +
+                           (output_height - y - 1) * dst_row_stride +
+                           (output_width - 1) * dst_pixel_stride;
+      for (dim_t x = 0; x < input_width; ++x) {
+        for (dim_t c = 0; c < depth; ++c) {
+          dst_ptr_row[c] = src_ptr_row[c];
+        }
+        dst_ptr_row -= depth;
+        src_ptr_row += depth;
+      }
+    }
+  }
+}
+
+// Rotates image for 270 degree.
+void Rotate270(dim_t batches, dim_t input_height, dim_t input_width,
+               dim_t depth, dim_t output_height, dim_t output_width,
+               const float* input_data, float* output_data) {
+  TFLITE_CHECK(input_data != nullptr);
+  TFLITE_CHECK(output_data != nullptr);
+
+  const int64_t pixel_stride = depth;
+  const int64_t src_row_stride = static_cast<int64_t>(input_width) * depth;
+  const int64_t dst_row_stride = static_cast<int64_t>(output_width) * depth;
+  const int64_t src_batch_stride = src_row_stride * input_height;
+  const int64_t dst_batch_stride = dst_row_stride * output_height;
+
+  // Iterate over batches to perform the following transformation:
+  // dst[output_height - x - 1][y] = src[y][x].
+  for (dim_t b = 0; b < batches; ++b) {
+    const float* src_data_ptr = input_data + b * src_batch_stride;
+    float* dst_data_ptr = output_data + b * dst_batch_stride;
+
+    for (dim_t y = 0; y < input_height; ++y) {
+      const float* src_ptr_row = src_data_ptr + y * src_row_stride;
+      for (dim_t x = 0; x < input_width; ++x) {
+        float* dst_ptr_row =
+            dst_data_ptr + (output_height - x - 1) * dst_row_stride;
+
+        const float* src_ptr_pixel = src_ptr_row + x * pixel_stride;
+        float* dst_pixel_ptr = dst_ptr_row + y * pixel_stride;
+
+        for (dim_t c = 0; c < depth; ++c) {
+          *dst_pixel_ptr++ = *src_ptr_pixel++;
+        }
+      }
+    }
+  }
+}
+
+// Performs generic rotation for arbitrary angle.
+void RotateGeneric(dim_t batches, dim_t input_height, dim_t input_width,
+                   dim_t depth, dim_t output_height, dim_t output_width,
+                   int angle, const float* input_data, float* output_data) {
+  TFLITE_CHECK(input_data != nullptr);
+  TFLITE_CHECK(output_data != nullptr);
+
+  const int64_t pixel_stride = depth;
+  const int64_t src_row_stride = static_cast<int64_t>(input_width) * depth;
+  const int64_t dst_row_stride = static_cast<int64_t>(output_width) * depth;
+  const int64_t src_batch_stride = src_row_stride * input_height;
+  const int64_t dst_batch_stride = dst_row_stride * output_height;
+
+  // Start off with dark image by initializing all pixels with zeros.
+  memset(output_data, 0,
+         static_cast<size_t>(batches) * output_width * output_height * depth *
+             sizeof(output_data[0]));
+
+  const float angle_rad = DegreesToRadians(angle);
+  const float cos_angle = std::cos(angle_rad);
+  const float sin_angle = std::sin(angle_rad);
+
+  const int64_t half_output_height = static_cast<int64_t>(output_height) / 2;
+  const int64_t half_output_width = static_cast<int64_t>(output_width) / 2;
+
+  // Iterate over batches to perform a rotation with arbitrary angle.
+  for (dim_t b = 0; b < batches; ++b) {
+    const float* src_data_ptr = input_data + b * src_batch_stride;
+    float* dst_data_ptr = output_data + b * dst_batch_stride;
+
+    for (int64_t y = -half_output_height; y < half_output_height; ++y) {
+      for (int64_t x = -half_output_width; x < half_output_width; ++x) {
+        const float x_transformed = cos_angle * x + sin_angle * y;
+        const float y_transformed = -sin_angle * x + cos_angle * y;
+
+        // Convert to integer by computing the next smaller integer number.
+        const int64_t x_transformed_integer =
+            static_cast<int64_t>(std::floor(x_transformed));
+        const int64_t y_transformed_integer =
+            static_cast<int64_t>(std::floor(y_transformed));
+
+        // Move into the coordinate system of input image.
+        const int64_t x_src_integer =
+            x_transformed_integer + static_cast<int64_t>(input_width) / 2;
+        const int64_t y_src_integer =
+            y_transformed_integer + static_cast<int64_t>(input_height) / 2;
+
+        // Calculate coordinates for interpolation.
+        const int64_t x0 = x_src_integer;
+        const int64_t x1 = x_src_integer + 1;
+        const int64_t y0 = y_src_integer;
+        const int64_t y1 = y_src_integer + 1;
+
+        // Skip further calculations if coordinates are out of bounds.
+        if (x0 < 0 || x0 >= input_width) continue;
+        if (x1 < 0 || x1 >= input_width) continue;
+        if (y0 < 0 || y0 >= input_height) continue;
+        if (y1 < 0 || y1 >= input_height) continue;
+
+        const float x_dist = x_transformed - x_transformed_integer;
+        const float y_dist = y_transformed - y_transformed_integer;
+        const float one_minus_x_dist = 1.0f - x_dist;
+        const float one_minus_y_dist = 1.0f - y_dist;
+
+        // Calculate rotated pixels for all channels.
+        const float* src_ptr_row0 = src_data_ptr + y0 * src_row_stride;
+        const float* src_ptr_row1 = src_data_ptr + y1 * src_row_stride;
+        float* dst_row_ptr =
+            dst_data_ptr + (y + half_output_height) * dst_row_stride;
+
+        const float* src_ptr_pixel00 = src_ptr_row0 + x0 * pixel_stride;
+        const float* src_ptr_pixel10 = src_ptr_row0 + x1 * pixel_stride;
+        const float* src_ptr_pixel01 = src_ptr_row1 + x0 * pixel_stride;
+        const float* src_ptr_pixel11 = src_ptr_row1 + x1 * pixel_stride;
+        float* dst_pixel_ptr =
+            dst_row_ptr + (x + half_output_width) * pixel_stride;
+
+        for (dim_t c = 0; c < depth; ++c) {
+          const float v00 = *src_ptr_pixel00++;
+          const float v01 = *src_ptr_pixel01++;
+          const float v10 = *src_ptr_pixel10++;
+          const float v11 = *src_ptr_pixel11++;
+
+          *dst_pixel_ptr++ =
+              (v10 * one_minus_y_dist + v11 * y_dist) * x_dist +
+              (v00 * one_minus_y_dist + v01 * y_dist) * one_minus_x_dist;
+        }
+      }
+    }
+  }
+}
+
+// Rotate given input with arbitrary angle. Works on `float` datatype.
+void ComputeRotate(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 2);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  MutableDataRef* output = outputs[0];
+
+  TFLITE_CHECK_EQ(img->Type(), etype_t::f32);
+  TFLITE_CHECK_EQ(output->Type(), etype_t::f32);
+
+  const float* img_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t img_num_batches = img->Dims()[0];
+  const dim_t img_height = img->Dims()[1];
+  const dim_t img_width = img->Dims()[2];
+  const dim_t img_num_channels = img->Dims()[3];
+
+  if (img_num_batches == 0 || img_height == 0 || img_width == 0) return;
+
+  const DataRef* angle = inputs[1];
+  TFLITE_CHECK_EQ(angle->NumElements(), 1);
+  TFLITE_CHECK_EQ(angle->Type(), etype_t::i32);
+  const int raw_angle = *reinterpret_cast<const int*>(angle->Data());
+  const int angle_data = ((raw_angle % 360) + 360) % 360;
+
+  // Resize output buffer for rotated image.
+  dim_t new_width = 0;
+  dim_t new_height = 0;
+  ComputeNewSize(img_width, img_height, angle_data, new_width, new_height);
+  output->Resize({img_num_batches, new_height, new_width, img_num_channels});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  // Perform rotation depending on angle.
+  if (angle_data == 90) {
+    Rotate90(img_num_batches, img_height, img_width, img_num_channels,
+             new_height, new_width, img_data, output_data);
+    return;
+  }
+
+  if (angle_data == 180) {
+    Rotate180(img_num_batches, img_height, img_width, img_num_channels,
+              new_height, new_width, img_data, output_data);
+    return;
+  }
+
+  if (angle_data == 270) {
+    Rotate270(img_num_batches, img_height, img_width, img_num_channels,
+              new_height, new_width, img_data, output_data);
+    return;
+  }
+
+  RotateGeneric(img_num_batches, img_height, img_width, img_num_channels,
+                new_height, new_width, angle_data, img_data, output_data);
+}
+
+}  // namespace
+
+const Algo* Impl_Rotate() {
+  static constexpr Algo kRotate = {&ComputeRotate, nullptr};
+  return &kRotate;
+}
+
+}  // namespace rotate
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/rotate.h

@@ -0,0 +1,39 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_ROTATE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_ROTATE_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace rotate {
+
+// Rotate
+//
+// Inputs: [img: float, angle: scalar<int>]
+// Ouputs: [img: float]
+//
+// Rotates the given image with arbitrary `angle`.
+// Overlaps with semantic of `tf.image.rotate` only for angles multiple to 90
+// degrees.
+//
+// https://www.tensorflow.org/api_docs/python/tf/image/rotate
+
+const algo::Algo* Impl_Rotate();
+
+}  // namespace rotate
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_ROTATE_H_

tensorflow/lite/experimental/ml_adjacent/algo/rotate_test.cc

@@ -0,0 +1,161 @@
+/* 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/experimental/ml_adjacent/algo/rotate.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace rotate {
+namespace {
+
+struct RotateTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const int angle;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class RotateTest : public ::testing::TestWithParam<RotateTestParams> {};
+
+TEST_P(RotateTest, FloatPixelType) {
+  constexpr float kAbsError = 0.01f;
+
+  const RotateTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Specify angle for rotation.
+  OwningVectorRef angle(etype_t::i32);
+  angle.Resize({1});
+  ASSERT_EQ(angle.Bytes(), sizeof(int));
+  std::memcpy(angle.Data(), &params.angle, angle.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Run image rotate custom call.
+  const Algo* rotate = Impl_Rotate();
+  rotate->process({&img, &angle}, {&output});
+
+  // Check rotate output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    RotateTests, RotateTest,
+    testing::ValuesIn({
+        // 4-D Tensor of shape [batch, height, width, channels] is used below.
+
+        RotateTestParams{/*img_dims=*/{1, 3, 3, 1},
+                         /*img_data=*/
+                         {11, 12, 13,  //
+                          21, 22, 23,  //
+                          31, 32, 33},
+                         /*angle=*/90,
+                         /*expected_data=*/
+                         {31, 21, 11,  //
+                          32, 22, 12,  //
+                          33, 23, 13},
+                         /*expected_shape=*/{1, 3, 3, 1}},
+
+        RotateTestParams{/*img_dims=*/{1, 3, 3, 1},
+                         /*img_data=*/
+                         {11, 12, 13,  //
+                          21, 22, 23,  //
+                          31, 32, 33},
+                         /*angle=*/180,
+                         /*expected_data=*/
+                         {33, 32, 31,  //
+                          23, 22, 21,  //
+                          13, 12, 11},
+                         /*expected_shape=*/{1, 3, 3, 1}},
+
+        RotateTestParams{/*img_dims=*/{1, 3, 3, 1},
+                         /*img_data=*/
+                         {11, 12, 13,  //
+                          21, 22, 23,  //
+                          31, 32, 33},
+                         /*angle=*/270,
+                         /*expected_data=*/
+                         {13, 23, 33,  //
+                          12, 22, 32,  //
+                          11, 21, 31},
+                         /*expected_shape=*/{1, 3, 3, 1}},
+
+        RotateTestParams{/*img_dims=*/{1, 8, 8, 1},
+                         /*img_data=*/
+                         {1, 1, 1, 1, 1, 1, 1, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 0, 0, 0, 0, 0, 0, 1,  //
+                          1, 1, 1, 1, 1, 1, 1, 1},
+                         /*angle=*/-45,
+                         /*expected_data=*/
+                         {
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.59f,
+                             0.83f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.54f, 0.00f,
+                             0.12f, 0.83f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.54f, 0.00f, 0.00f,
+                             0.00f, 0.12f, 0.83f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.54f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.12f, 0.83f, 0.00f, 0.00f,  //
+                             0.00f, 0.78f, 0.00f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.00f, 0.23f, 0.97f, 0.00f,  //
+                             0.00f, 0.00f, 0.54f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.12f, 0.83f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.54f, 0.00f, 0.00f,
+                             0.00f, 0.12f, 0.83f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.54f, 0.00f,
+                             0.12f, 0.83f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.59f,
+                             0.83f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,
+                             0.00f, 0.00f, 0.00f, 0.00f, 0.00f, 0.00f,  //
+                         },
+                         /*expected_shape=*/{1, 12, 12, 1}},
+    }));
+
+}  // namespace
+}  // namespace rotate
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/yuv_to_rgb.cc

@@ -0,0 +1,75 @@
+/* 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/experimental/ml_adjacent/algo/image_utils.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace ml_adj {
+namespace yuv_to_rgb {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::DataRef;
+using ::ml_adj::data::MutableDataRef;
+
+constexpr float kYuv2RgbKernel[] = {1.0f,         0.0f,
+                                    1.13988303f,  //
+                                    1.0f,         -0.394642334f,
+                                    -0.58062185f,  //
+                                    1.0f,         2.03206185f,   0.0f};
+constexpr int kYuv2RgbKernelDim =
+    sizeof(kYuv2RgbKernel) / sizeof(kYuv2RgbKernel[0]);
+
+// Converts YUV to RGB. Supports `float` datatype only.
+void ComputeYuvToRgb(const InputPack& inputs, const OutputPack& outputs) {
+  TFLITE_CHECK_EQ(inputs.size(), 1);
+  TFLITE_CHECK_EQ(outputs.size(), 1);
+
+  // Extract input image data.
+  const DataRef* img = inputs[0];
+  TFLITE_CHECK_EQ(img->Dims().size(), 4);
+  MutableDataRef* output = outputs[0];
+
+  TFLITE_CHECK_EQ(img->Type(), etype_t::f32);
+  TFLITE_CHECK_EQ(output->Type(), etype_t::f32);
+
+  const float* input_data = reinterpret_cast<const float*>(img->Data());
+  const dim_t batches = img->Dims()[0];
+  const dim_t height = img->Dims()[1];
+  const dim_t width = img->Dims()[2];
+  const dim_t channels = img->Dims()[3];
+  TFLITE_CHECK_EQ(channels, 3);
+
+  if (batches == 0 || height == 0 || width == 0) return;
+
+  // Resize output buffer.
+  output->Resize({batches, height, width, channels});
+  float* output_data = reinterpret_cast<float*>(output->Data());
+
+  ConvertColorSpace(batches, height, width, input_data, output_data,
+                    &kYuv2RgbKernel[0], kYuv2RgbKernelDim);
+}
+
+}  // namespace
+
+const Algo* Impl_YuvToRgb() {
+  static constexpr Algo kYuvToRgb = {&ComputeYuvToRgb, nullptr};
+  return &kYuvToRgb;
+}
+
+}  // namespace yuv_to_rgb
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/algo/yuv_to_rgb.h

@@ -0,0 +1,38 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_ALGO_YUV_TO_RGB_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_YUV_TO_RGB_H_
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace yuv_to_rgb {
+
+// YUV to RGB conversion
+//
+// Inputs: [img: float]
+// Ouputs: [img: float]
+//
+// Coverts the given 3-channel YUV image to 3-channel RGB image.
+// Mimics semantic of `tf.image.yuv_to_rgb.
+
+// https://www.tensorflow.org/api_docs/python/tf/image/yuv_to_rgb
+
+const algo::Algo* Impl_YuvToRgb();
+
+}  // namespace yuv_to_rgb
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_ALGO_YUV_TO_RGB_H_

tensorflow/lite/experimental/ml_adjacent/algo/yuv_to_rgb_test.cc

@@ -0,0 +1,119 @@
+/* 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/experimental/ml_adjacent/algo/yuv_to_rgb.h"
+
+#include <cstring>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::data::OwningVectorRef;
+
+namespace ml_adj {
+namespace yuv_to_rgb {
+namespace {
+
+struct YuvToRgbTestParams {
+  const std::vector<dim_t> img_dims;
+  const std::vector<float> img_data;
+  const std::vector<float> expected_data;
+  const std::vector<dim_t> expected_shape;
+};
+
+class YuvToRgbTest : public ::testing::TestWithParam<YuvToRgbTestParams> {};
+
+TEST_P(YuvToRgbTest, FloatPixelType) {
+  constexpr float kAbsError = 0.1f;
+  const YuvToRgbTestParams& params = GetParam();
+
+  // Image input.
+  OwningVectorRef img(etype_t::f32);
+  img.Resize(dims_t(params.img_dims));
+  ASSERT_EQ(img.Bytes(), params.img_data.size() * sizeof(float));
+  std::memcpy(img.Data(), params.img_data.data(), img.Bytes());
+
+  // Empty output image.
+  OwningVectorRef output(etype_t::f32);
+
+  // Convert YUV to RGB image.
+  const Algo* yuv_to_rgb = Impl_YuvToRgb();
+  yuv_to_rgb->process({&img}, {&output});
+
+  // Check resize output.
+  ASSERT_EQ(output.Bytes(), params.expected_data.size() * sizeof(float));
+  ASSERT_EQ(output.Dims(), params.expected_shape);
+
+  const float* out_data = reinterpret_cast<float*>(output.Data());
+  for (int i = 0; i < output.NumElements(); ++i) {
+    EXPECT_NEAR(out_data[i], params.expected_data[i], kAbsError)
+        << "out_data[" << i << "] = " << out_data[i] << ", expected_data[" << i
+        << "] = " << params.expected_data[i];
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    YuvToRgbTests, YuvToRgbTest,
+    testing::ValuesIn({
+        YuvToRgbTestParams{/*img_dims=*/{1, 3, 2, 3},
+                           /*img_data=*/
+                           {
+                               92.5f,
+                               58.3f,
+                               -71.5f,
+                               93.5f,
+                               58.3f,
+                               -71.5f,  //
+                               102.5f,
+                               58.3f,
+                               -71.5f,
+                               103.5f,
+                               58.3f,
+                               -71.5f,  //
+                               112.5f,
+                               58.3f,
+                               -71.5f,
+                               113.5f,
+                               58.3f,
+                               -71.5f,
+                           },
+                           /*expected_data=*/
+                           {11, 111, 211, 12, 112, 212,  //
+                            21, 121, 221, 22, 122, 222,  //
+                            31, 131, 231, 32, 132, 232},
+                           /*expected_shape=*/{1, 3, 2, 3}},
+        YuvToRgbTestParams{/*img_dims=*/{2, 3, 2, 3},
+                           /*img_data=*/
+                           {92.5f,  58.3f, -71.5f, 93.5f,  58.3f, -71.5f,  //
+                            102.5f, 58.3f, -71.5f, 103.5f, 58.3f, -71.5f,  //
+                            112.5f, 58.3f, -71.5f, 113.5f, 58.3f, -71.5f,  //
+                            92.5f,  58.3f, -71.5f, 93.5f,  58.3f, -71.5f,  //
+                            102.5f, 58.3f, -71.5f, 103.5f, 58.3f, -71.5f,  //
+                            112.5f, 58.3f, -71.5f, 113.5f, 58.3f, -71.5f},
+                           /*expected_data=*/
+                           {11, 111, 211, 12, 112, 212,  //
+                            21, 121, 221, 22, 122, 222,  //
+                            31, 131, 231, 32, 132, 232,  //
+                            11, 111, 211, 12, 112, 212,  //
+                            21, 121, 221, 22, 122, 222,  //
+                            31, 131, 231, 32, 132, 232},
+                           /*expected_shape=*/{2, 3, 2, 3}},
+    }));
+
+}  // namespace
+}  // namespace yuv_to_rgb
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/data/BUILD

@@ -0,0 +1,27 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "owning_vector_ref",
+    srcs = ["owning_vector_ref.cc"],
+    hdrs = [
+        "owning_vector_ref.h",
+    ],
+    deps = ["//tensorflow/lite/experimental/ml_adjacent:lib"],
+)
+
+cc_test(
+    name = "owning_vector_ref_test",
+    srcs = ["owning_vector_ref_test.cc"],
+    deps = [
+        ":owning_vector_ref",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.cc

@@ -0,0 +1,54 @@
+/* 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/experimental/ml_adjacent/data/owning_vector_ref.h"
+
+#include <cstddef>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace data {
+
+void OwningVectorRef::Resize(dims_t&& dims) {
+  dims_ = dims;
+
+  num_elements_ = 0;
+  // TODO(b/292143456) Add a helper for this.
+  for (dim_t d : dims_) {
+    if (d <= 0) {
+      break;
+    }
+    if (num_elements_ == 0) {
+      num_elements_ = d;
+    } else {
+      num_elements_ *= d;
+    }
+  }
+
+  raw_data_buffer_.resize(num_elements_ * TypeWidth(Type()));
+}
+
+const void* OwningVectorRef::Data() const { return raw_data_buffer_.data(); }
+
+void* OwningVectorRef::Data() { return raw_data_buffer_.data(); }
+
+ind_t OwningVectorRef::NumElements() const { return num_elements_; }
+
+size_t OwningVectorRef::Bytes() const {
+  return NumElements() * TypeWidth(Type());
+}
+
+}  // namespace data
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref.h

@@ -0,0 +1,58 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_DATA_OWNING_VECTOR_REF_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_DATA_OWNING_VECTOR_REF_H_
+#include <vector>
+
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace data {
+
+// A MutableDataRef implemenation that manages its buffer through underlying
+// vector which it owns.
+class OwningVectorRef : public MutableDataRef {
+ public:
+  explicit OwningVectorRef(etype_t type) : MutableDataRef(type) {}
+
+  OwningVectorRef(const OwningVectorRef&) = delete;
+  OwningVectorRef(OwningVectorRef&&) = delete;
+  OwningVectorRef& operator=(const OwningVectorRef&) = delete;
+  OwningVectorRef& operator=(OwningVectorRef&&) = delete;
+
+  // Resizes the underlying vector to prod(dims) * type width.
+  void Resize(dims_t&& dims) override;
+
+  // Gets read-only pointer to vector's buffer.
+  const void* Data() const override;
+
+  // Gets a write-read pointer to vector's buffer.
+  void* Data() override;
+
+  ind_t NumElements() const override;
+
+  size_t Bytes() const override;
+
+  ~OwningVectorRef() override = default;
+
+ private:
+  std::vector<char> raw_data_buffer_;
+  ind_t num_elements_ = 0;
+};
+
+}  // namespace data
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_DATA_OWNING_VECTOR_REF_H_

tensorflow/lite/experimental/ml_adjacent/data/owning_vector_ref_test.cc

@@ -0,0 +1,89 @@
+/* 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/experimental/ml_adjacent/data/owning_vector_ref.h"
+
+#include <algorithm>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj {
+namespace data {
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::IsEmpty;
+
+TEST(OwningVectorRefTest, ConstructFloat32) {
+  OwningVectorRef t(etype_t::f32);
+
+  EXPECT_EQ(t.Type(), etype_t::f32);
+  EXPECT_EQ(t.NumElements(), 0);
+  EXPECT_EQ(t.Bytes(), 0);
+  EXPECT_THAT(t.Dims(), IsEmpty());
+}
+
+TEST(OwningVectorRefTest, ResizeFromEmptyFloat32) {
+  OwningVectorRef t(etype_t::f32);
+  t.Resize({2, 2});
+
+  EXPECT_THAT(t.Dims(), ElementsAreArray<dim_t>({2, 2}));
+  EXPECT_EQ(t.NumElements(), 4);
+  ASSERT_EQ(t.Bytes(), 4 * sizeof(float));
+
+  // Check buffer is correct.
+  float* write_f_start = reinterpret_cast<float*>(t.Data());
+  float* write_f_end = write_f_start + t.NumElements();
+  std::fill(write_f_start, write_f_end, 0.5f);
+
+  const float* read_f_start = reinterpret_cast<const float*>(t.Data());
+  for (int i = 0; i < t.NumElements(); ++i) {
+    EXPECT_EQ(read_f_start[i], 0.5f);
+  }
+}
+
+TEST(OwningVectorRefTest, ResizeDownFloat32) {
+  OwningVectorRef t(etype_t::f32);
+  t.Resize({2, 2});
+
+  float* write_f_start = reinterpret_cast<float*>(t.Data());
+  float* write_f_end = write_f_start + t.NumElements();
+  std::fill(write_f_start, write_f_end, 0.5f);
+
+  t.Resize({3});
+  ASSERT_THAT(t.Dims(), ElementsAreArray<dim_t>({3}));
+  EXPECT_EQ(t.NumElements(), 3);
+  ASSERT_EQ(t.Bytes(), 3 * sizeof(float));
+
+  const float* read_f_start = reinterpret_cast<const float*>(t.Data());
+  for (int i = 0; i < t.NumElements(); ++i) {
+    EXPECT_EQ(read_f_start[i], 0.5f);
+  }
+}
+
+TEST(OwningVectorRefTest, IgnoresDimsForNumElementsAfterFirstNonPositive) {
+  OwningVectorRef t(etype_t::f32);
+  t.Resize({3, 0, 0, 2});
+
+  EXPECT_EQ(t.Type(), etype_t::f32);
+  EXPECT_EQ(t.NumElements(), 3);
+  EXPECT_EQ(t.Bytes(), 3 * sizeof(float));
+  EXPECT_THAT(t.Dims(), ElementsAreArray<dim_t>({3, 0, 0, 2}));
+}
+
+}  // namespace
+}  // namespace data
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/lib.h

@@ -0,0 +1,139 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_LIB_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_LIB_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+namespace ml_adj {
+
+/// Standard Types ///
+
+// Length of axis.
+typedef uint32_t dim_t;
+
+// Dimensions of data.
+typedef std::vector<dim_t> dims_t;
+
+// 1d index into data.
+typedef uint64_t ind_t;
+
+// Integral type of data element.
+enum etype_t : uint8_t {
+  i32 = 0,
+  f32 = 1,
+  f64 = 2,
+};
+
+// Size in bytes of data element.
+typedef uint8_t width_t;
+
+namespace data {
+
+// Get size (int bytes) of single element of type.
+inline width_t TypeWidth(etype_t type) {
+  switch (type) {
+    case etype_t::i32:
+      return sizeof(int32_t);
+    case etype_t::f32:
+      return sizeof(float);
+    case etype_t::f64:
+      return sizeof(double);
+  }
+}
+
+/// Input/Output Wrapper for Algos ///
+
+// Encapsulates a input or output to an algorithm. Management of
+// buffers is to be implemented outside of algorithms.
+
+// Read only wrapper.
+class DataRef {
+ public:
+  explicit DataRef(etype_t type) : element_type_(type) {}
+
+  DataRef(const DataRef&) = delete;
+  DataRef(DataRef&&) = delete;
+  DataRef& operator=(const DataRef&) = delete;
+  DataRef& operator=(DataRef&&) = delete;
+
+  // Read only buffer, allocated to be of size == Bytes().
+  virtual const void* Data() const = 0;
+
+  // Number of elements currently allocated.
+  virtual ind_t NumElements() const = 0;
+
+  // Size of buffer.
+  virtual size_t Bytes() const = 0;
+
+  // Type of elements.
+  etype_t Type() const { return element_type_; }
+
+  // Implicit dimensions of buffer.
+  const dims_t& Dims() const { return dims_; }
+
+  virtual ~DataRef() = default;
+
+ protected:
+  dims_t dims_;
+  // Data can be reshaped but not change types.
+  const etype_t element_type_;
+};
+
+// Read/write wrapper which can be resized.
+class MutableDataRef : public DataRef {
+ public:
+  using DataRef::Data;
+
+  explicit MutableDataRef(etype_t type) : DataRef(type) {}
+
+  // Takes ownership of dims_t. Implementations must set bytes and
+  // `num_elements_` field.
+  virtual void Resize(dims_t&& dims) = 0;
+
+  // Write buffer, allocated to be of size == Bytes().
+  virtual void* Data() = 0;
+};
+
+}  // namespace data
+
+namespace algo {
+
+/// Function Interface for Operations on DataRefs ///
+
+// Inputs to algorithm.
+typedef std::vector<data::DataRef*> InputPack;
+// Outputs to algorithm.
+typedef std::vector<data::MutableDataRef*> OutputPack;
+
+// Generic algorithm, computes outputs via inputs.
+typedef void ComputeFunc(const InputPack& inputs, const OutputPack& outputs);
+
+// Optional hook to compute output shapes when they are non data-dependent.
+// This is a place-holder for now.
+// TODO(b/292143456) Figure out what the signature of this should be.
+typedef size_t ShapeFunc();
+
+struct Algo {
+  ComputeFunc* process = nullptr;
+  ShapeFunc* output_size = nullptr;
+};
+
+}  // namespace algo
+}  // namespace ml_adj
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_LIB_H_

tensorflow/lite/experimental/ml_adjacent/tflite/BUILD

@@ -0,0 +1,67 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tfl_tensor_ref",
+    srcs = ["tfl_tensor_ref.cc"],
+    hdrs = ["tfl_tensor_ref.h"],
+    deps = [
+        "//tensorflow/lite:array",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels/internal:compatibility",
+    ],
+)
+
+cc_test(
+    name = "tfl_tensor_ref_test",
+    srcs = ["tfl_tensor_ref_test.cc"],
+    deps = [
+        ":tfl_tensor_ref",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "extern_call",
+    srcs = ["extern_call.cc"],
+    hdrs = ["extern_call.h"],
+    deps = [
+        ":tfl_tensor_ref",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/experimental/ml_adjacent:lib",
+        "//tensorflow/lite/experimental/ml_adjacent/algo:crop",
+        "//tensorflow/lite/experimental/ml_adjacent/algo:resize",
+        "//tensorflow/lite/kernels:kernel_util",
+    ],
+)
+
+cc_test(
+    name = "extern_call_test",
+    srcs = ["extern_call_test.cc"],
+    deps = [
+        ":extern_call",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:test_util",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/ml_adjacent/tflite/extern_call.cc

@@ -0,0 +1,106 @@
+/* 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/experimental/ml_adjacent/tflite/extern_call.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/ml_adjacent/algo/crop.h"
+#include "tensorflow/lite/experimental/ml_adjacent/algo/resize.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/experimental/ml_adjacent/tflite/tfl_tensor_ref.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace extern_call {
+namespace {
+
+using ::ml_adj::algo::Algo;
+using ::ml_adj::algo::InputPack;
+using ::ml_adj::algo::OutputPack;
+using ::ml_adj::data::MutableTflTensorRef;
+using ::ml_adj::data::TflTensorRef;
+
+// UniquePtr wrapper around vectors that hold the inputs/outputs to
+// library's `Algo`s.
+template <typename PackType>
+struct PackDeleter {
+  void operator()(PackType* pack) {
+    if (pack == nullptr) return;
+    for (auto* d : *pack) {
+      if (d == nullptr) continue;
+      delete d;
+    }
+    delete pack;
+  }
+};
+template <typename PackType>
+using UniquePack = std::unique_ptr<PackType, PackDeleter<PackType>>;
+
+constexpr uint8_t kNumFuncs = 2;
+static const Algo* const kReg[kNumFuncs] = {ml_adj::crop::Impl_CenterCrop(),
+                                            ml_adj::resize::Impl_Resize()};
+
+TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
+  for (int i = 0; i < NumOutputs(node); ++i) {
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i, &output));
+    SetTensorToDynamic(output);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
+  UniquePack<InputPack> lib_inputs(new InputPack());
+  for (int i = 0; i < NumInputs(node); ++i) {
+    const TfLiteTensor* input;
+    TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &input));
+
+    lib_inputs->push_back(new TflTensorRef(input));
+  }
+
+  UniquePack<OutputPack> lib_outputs(new OutputPack());
+  for (int i = 0; i < NumOutputs(node); ++i) {
+    TfLiteTensor* output;
+    TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, i, &output));
+
+    lib_outputs->push_back(new MutableTflTensorRef(output, context));
+  }
+
+  TF_LITE_ENSURE_EQ(context, node->custom_initial_data_size,
+                    sizeof(ExternCallOptions));
+  const auto* const options =
+      reinterpret_cast<const ExternCallOptions*>(node->custom_initial_data);
+  TF_LITE_ENSURE(context,
+                 options->func_id >= 0 && options->func_id < kNumFuncs);
+
+  const Algo* const algo = kReg[options->func_id];
+
+  algo->process(*lib_inputs, *lib_outputs);
+
+  return kTfLiteOk;
+}
+}  // namespace
+
+TfLiteRegistration* Register_EXTERN_CALL() {
+  static TfLiteRegistration r = {nullptr, nullptr, Prepare, Eval};
+  return &r;
+}
+
+}  // namespace extern_call
+}  // namespace tflite

tensorflow/lite/experimental/ml_adjacent/tflite/extern_call.h

@@ -0,0 +1,37 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_TFLITE_EXTERN_CALL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_TFLITE_EXTERN_CALL_H_
+
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite::extern_call {
+
+// Compile time options passed to this kernel at runtime.
+struct ExternCallOptions {
+  // A single custom op is used to represent a call to an arbitrary function
+  // in the library. The function that is called is encoded in `func_id`.
+  // Because of compiler op def, these will be encded at compile time
+  // as `char[]` and will serialize `uint8_t`, so we match this type.
+  uint8_t func_id;
+};
+
+TfLiteRegistration* Register_EXTERN_CALL();
+
+}  // namespace tflite::extern_call
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_TFLITE_EXTERN_CALL_H_

tensorflow/lite/experimental/ml_adjacent/tflite/extern_call_test.cc

@@ -0,0 +1,90 @@
+/* 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/experimental/ml_adjacent/tflite/extern_call.h"
+
+#include <cstdint>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+class ExternCallModel : public SingleOpModel {
+ public:
+  ExternCallModel(std::vector<TensorData> inputs,
+                  std::vector<TensorData> outputs, uint8_t func_id) {
+    std::vector<std::vector<int>> input_shapes;
+    for (const auto& data : inputs) {
+      input_ids_.push_back(AddInput(data));
+      input_shapes.push_back(GetShape(input_ids_.back()));
+    }
+    for (const auto& data : outputs) {
+      output_ids_.push_back(AddOutput(data));
+    }
+    SetCustomOp("ExternCall", {func_id}, extern_call::Register_EXTERN_CALL);
+    BuildInterpreter(input_shapes);
+  }
+
+  const TfLiteTensor* Output(int output_id) {
+    return interpreter_->tensor(output_ids_[output_id]);
+  }
+
+ private:
+  std::vector<int> input_ids_;
+  std::vector<int> output_ids_;
+};
+
+namespace {
+
+// This custom op is a simple wrapper and most of the meaningful logic
+// with be tested in the contained `Algo` implementations. It is
+// sufficient to just test that calling these `Algo` succeeds.
+
+// TODO(b/290283768) If the complexity if this custom op grows consider
+// devising a mechanism allowing for a "mock" `Algo` to call from these tests.
+
+TEST(ExternCallTest, CropFunc) {
+  std::vector<TensorData> inputs = {{TensorType_FLOAT32, {1, 5, 5, 1}},
+                                    {TensorType_FLOAT64, {}}};
+  std::vector<TensorData> output = {{TensorType_FLOAT32, {}}};
+
+  ExternCallModel model(inputs, output, 0);
+  model.PopulateTensor<double>(1, {0.5});
+
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  ASSERT_NE(model.Output(0), nullptr);
+  ASSERT_THAT(model.Output(0), DimsAre({1, 3, 3, 1}));
+}
+
+TEST(ExternCallTest, ResizeTest) {
+  std::vector<TensorData> inputs = {{TensorType_FLOAT32, {1, 5, 5, 1}},
+                                    {TensorType_UINT32, {2}}};
+  std::vector<TensorData> output = {{TensorType_FLOAT32, {}}};
+
+  ExternCallModel model(inputs, output, 1);
+  model.PopulateTensor<uint32_t>(1, {3, 3});
+
+  ASSERT_EQ(model.Invoke(), kTfLiteOk);
+  ASSERT_NE(model.Output(0), nullptr);
+  ASSERT_THAT(model.Output(0), DimsAre({1, 3, 3, 1}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/experimental/ml_adjacent/tflite/tfl_tensor_ref.cc

@@ -0,0 +1,93 @@
+/* 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/experimental/ml_adjacent/tflite/tfl_tensor_ref.h"
+
+#include <cstddef>
+#include <vector>
+
+#include "tensorflow/lite/array.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace ml_adj {
+namespace data {
+
+using ::tflite::BuildTfLiteArray;
+using ::tflite::TfLiteArrayUniquePtr;
+using ::tflite::TfLiteTypeGetSize;
+
+namespace {
+etype_t TflToLibType(const TfLiteType tfl_type) {
+  switch (tfl_type) {
+    case kTfLiteFloat32:
+      return etype_t::f32;
+    case kTfLiteInt32:
+      return etype_t::i32;
+    case kTfLiteFloat64:
+      return etype_t::f64;
+    default:
+      return etype_t::i32;
+  }
+}
+}  // namespace
+
+TflTensorRef::TflTensorRef(const TfLiteTensor* tfl_tensor)
+    : DataRef(TflToLibType(tfl_tensor->type)), tfl_tensor_(tfl_tensor) {
+  dims_.assign(tfl_tensor->dims->data,
+               tfl_tensor->dims->data + tfl_tensor->dims->size);
+}
+
+const void* TflTensorRef::Data() const { return tfl_tensor_->data.data; }
+
+ind_t TflTensorRef::NumElements() const {
+  return tfl_tensor_->bytes / TfLiteTypeGetSize(tfl_tensor_->type);
+}
+
+size_t TflTensorRef::Bytes() const { return tfl_tensor_->bytes; }
+
+MutableTflTensorRef::MutableTflTensorRef(TfLiteTensor* tfl_tensor,
+                                         TfLiteContext* tfl_ctx)
+    : MutableDataRef(TflToLibType(tfl_tensor->type)),
+      tfl_tensor_(tfl_tensor),
+      tfl_ctx_(tfl_ctx) {
+  dims_.assign(tfl_tensor->dims->data,
+               tfl_tensor->dims->data + tfl_tensor->dims->size);
+}
+
+void MutableTflTensorRef::Resize(dims_t&& dims) {
+  TfLiteArrayUniquePtr<int> arr =
+      BuildTfLiteArray(std::vector<int>(dims.begin(), dims.end()));
+
+  TFLITE_CHECK_EQ(tfl_ctx_->ResizeTensor(tfl_ctx_, tfl_tensor_, arr.release()),
+                  kTfLiteOk);
+
+  dims_ = dims;
+}
+
+const void* MutableTflTensorRef::Data() const { return tfl_tensor_->data.data; }
+
+ind_t MutableTflTensorRef::NumElements() const {
+  return tfl_tensor_->bytes / TfLiteTypeGetSize(tfl_tensor_->type);
+}
+
+size_t MutableTflTensorRef::Bytes() const { return tfl_tensor_->bytes; }
+
+void* MutableTflTensorRef::Data() { return tfl_tensor_->data.data; }
+
+}  // namespace data
+}  // namespace ml_adj

tensorflow/lite/experimental/ml_adjacent/tflite/tfl_tensor_ref.h

@@ -0,0 +1,74 @@
+/* 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_EXPERIMENTAL_ML_ADJACENT_TFLITE_TFL_TENSOR_REF_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_TFLITE_TFL_TENSOR_REF_H_
+
+#include <cstddef>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+
+namespace ml_adj::data {
+
+// Immutable wrapper of `TfLiteTensor`.
+class TflTensorRef : public DataRef {
+ public:
+  explicit TflTensorRef(const TfLiteTensor* tfl_tensor);
+
+  // TODO(b/290283768) Implement copy and move semantics.
+  TflTensorRef(const TflTensorRef&) = delete;
+  TflTensorRef(TflTensorRef&&) = delete;
+  TflTensorRef& operator=(const TflTensorRef&) = delete;
+  TflTensorRef& operator=(TflTensorRef&&) = delete;
+
+  const void* Data() const override;
+
+  ind_t NumElements() const override;
+
+  size_t Bytes() const override;
+
+ private:
+  const TfLiteTensor* const tfl_tensor_;
+};
+
+// Mutable wrapper of `TfLiteTensor`.
+class MutableTflTensorRef : public MutableDataRef {
+ public:
+  MutableTflTensorRef(TfLiteTensor* tfl_tensor, TfLiteContext* tfl_ctx);
+
+  // TODO(b/290283768) Implement copy and move semantics.
+  MutableTflTensorRef(const MutableTflTensorRef&) = delete;
+  MutableTflTensorRef(MutableTflTensorRef&&) = delete;
+  MutableTflTensorRef& operator=(const MutableTflTensorRef&) = delete;
+  MutableTflTensorRef& operator=(MutableTflTensorRef&&) = delete;
+
+  const void* Data() const override;
+
+  ind_t NumElements() const override;
+
+  size_t Bytes() const override;
+
+  void Resize(dims_t&& dims) override;
+
+  void* Data() override;
+
+ private:
+  TfLiteTensor* const tfl_tensor_;
+  TfLiteContext* tfl_ctx_;
+};
+
+}  // namespace ml_adj::data
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ML_ADJACENT_TFLITE_TFL_TENSOR_REF_H_

tensorflow/lite/experimental/ml_adjacent/tflite/tfl_tensor_ref_test.cc

@@ -0,0 +1,156 @@
+/* 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/experimental/ml_adjacent/tflite/tfl_tensor_ref.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/experimental/ml_adjacent/lib.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace ml_adj {
+namespace data {
+namespace {
+
+using ::testing::Each;
+using ::tflite::BuildTfLiteTensor;
+using ::tflite::DimsAre;
+using ::tflite::NumElements;
+using ::tflite::TensorUniquePtr;
+
+// Mock implementation of `ResizeTensor`.
+TfLiteStatus SimpleResizeTensor(TfLiteContext*, TfLiteTensor* tensor,
+                                TfLiteIntArray* new_size) {
+  TFLITE_CHECK(tensor->type == kTfLiteFloat32);
+  size_t num_bytes = NumElements(new_size) * sizeof(float);
+  TF_LITE_ENSURE_STATUS(TfLiteTensorRealloc(num_bytes, tensor));
+  if (tensor->dims != nullptr) {
+    TfLiteIntArrayFree(tensor->dims);
+  }
+  tensor->dims = new_size;
+  return kTfLiteOk;
+}
+
+std::unique_ptr<TfLiteContext> MakeSimpleContext() {
+  auto ctx = std::make_unique<TfLiteContext>();
+  ctx->ResizeTensor = SimpleResizeTensor;
+  return ctx;
+}
+
+TEST(ImmutableTensorRefTest, ConstructsAndManifestsTensorData) {
+  TensorUniquePtr tfl_tensor =
+      BuildTfLiteTensor(kTfLiteFloat32, {2, 2}, kTfLiteDynamic);
+  std::fill(tfl_tensor->data.f, tfl_tensor->data.f + 4, 2.0f);
+  TflTensorRef ref(tfl_tensor.get());
+
+  ASSERT_EQ(ref.Type(), etype_t::f32);
+  ASSERT_EQ(ref.Dims(), (dims_t{2, 2}));
+  ASSERT_EQ(ref.Bytes(), 4 * sizeof(float));
+
+  absl::Span<const float> data(reinterpret_cast<const float*>(ref.Data()), 4);
+  EXPECT_THAT(data, Each(2.0f));
+}
+
+TEST(MutableTensorRefTest, ConstructsAndManifestsTensorData) {
+  TensorUniquePtr tfl_tensor =
+      BuildTfLiteTensor(kTfLiteFloat32, {2, 2}, kTfLiteDynamic);
+  std::fill(tfl_tensor->data.f, tfl_tensor->data.f + 4, 2.0f);
+  MutableTflTensorRef ref(tfl_tensor.get(), nullptr);
+
+  ASSERT_EQ(ref.Type(), etype_t::f32);
+  ASSERT_EQ(ref.Dims(), (dims_t{2, 2}));
+  ASSERT_EQ(ref.Bytes(), 4 * sizeof(float));
+
+  absl::Span<const float> data(reinterpret_cast<const float*>(ref.Data()), 4);
+  EXPECT_THAT(data, Each(2.0f));
+}
+
+TEST(MutableTensorRefTest, TensorRefWritesDataToTensor) {
+  TensorUniquePtr tfl_tensor =
+      BuildTfLiteTensor(kTfLiteFloat32, {3, 3}, kTfLiteDynamic);
+  MutableTflTensorRef ref(tfl_tensor.get(), nullptr);
+
+  ASSERT_EQ(ref.Type(), etype_t::f32);
+  ASSERT_EQ(ref.Dims(), (dims_t{3, 3}));
+  ASSERT_EQ(ref.Bytes(), 9 * sizeof(float));
+
+  absl::Span<float> data(reinterpret_cast<float*>(ref.Data()), 9);
+  std::fill(data.begin(), data.end(), 3.0f);
+
+  EXPECT_THAT(absl::Span<const float>(tfl_tensor->data.f, 9), Each(3.0f));
+}
+
+TEST(MutableTensorRefTest, ResizeIncreaseSize) {
+  TensorUniquePtr tfl_tensor =
+      BuildTfLiteTensor(kTfLiteFloat32, {2, 3}, kTfLiteDynamic);
+  std::unique_ptr<TfLiteContext> ctx = MakeSimpleContext();
+  MutableTflTensorRef ref(tfl_tensor.get(), ctx.get());
+
+  ASSERT_EQ(ref.Type(), etype_t::f32);
+  ASSERT_EQ(ref.Dims(), (dims_t{2, 3}));
+  ASSERT_EQ(ref.Bytes(), 6 * sizeof(float));
+
+  ref.Resize({3, 3});
+  ASSERT_EQ(ref.Dims(), (dims_t{3, 3}));
+  ASSERT_EQ(ref.Bytes(), 9 * sizeof(float));
+
+  // Sanitizers will check buffer is correct size.
+  absl::Span<float> ref_data(reinterpret_cast<float*>(ref.Data()), 9);
+
+  // Check underlying tensor is also resized.
+  ASSERT_THAT(tfl_tensor.get(), DimsAre({3, 3}));
+  ASSERT_EQ(tfl_tensor->bytes, ref.Bytes());
+
+  // Check share same buffer.
+  ASSERT_EQ(ref.Data(), tfl_tensor->data.data);
+}
+
+TEST(MutableTensorRefTest, ResizeDecreasesSize) {
+  TensorUniquePtr tfl_tensor =
+      BuildTfLiteTensor(kTfLiteFloat32, {2, 3}, kTfLiteDynamic);
+  std::unique_ptr<TfLiteContext> ctx = MakeSimpleContext();
+  MutableTflTensorRef ref(tfl_tensor.get(), ctx.get());
+
+  ASSERT_EQ(ref.Type(), etype_t::f32);
+  ASSERT_EQ(ref.Dims(), (dims_t{2, 3}));
+  ASSERT_EQ(ref.Bytes(), 6 * sizeof(float));
+
+  ref.Resize({2, 2});
+  ASSERT_EQ(ref.Dims(), (dims_t{2, 2}));
+  ASSERT_EQ(ref.Bytes(), 4 * sizeof(float));
+
+  // Sanitizers will check buffer is correct size.
+  absl::Span<float> ref_data(reinterpret_cast<float*>(ref.Data()), 4);
+
+  // Check underlying tensor is also resized.
+  ASSERT_THAT(tfl_tensor.get(), DimsAre({2, 2}));
+  ASSERT_EQ(tfl_tensor->bytes, ref.Bytes());
+
+  // Check share same buffer.
+  ASSERT_EQ(ref.Data(), tfl_tensor->data.data);
+}
+
+}  // namespace
+}  // namespace data
+}  // namespace ml_adj