Port reference squared difference kernel to TFLM (#1115)

* Copy squared_diffence* from TFLite verbatim

* Remove TFLite specific contents and flatten namespaces

* Complete porting squared difference to micro

* Updated based on comments
Co-authored-by: Nmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tensorflow/lite/micro/kernels/BUILD

@@ -231,6 +231,7 @@ tflm_kernel_cc_library(
         "space_to_depth.cc",
         "split.cc",
         "split_v.cc",
+        "squared_difference.cc",
         "squeeze.cc",
         "strided_slice.cc",
         "sub.cc",
@@ -1198,6 +1199,19 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "squared_difference_test",
+    srcs = [
+        "squared_difference_test.cc",
+    ],
+    deps = [
+        ":kernel_runner",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/micro:test_helpers",
+        "//tensorflow/lite/micro/testing:micro_test",
+    ],
+)
+
 cc_test(
     name = "squeeze_test",
     srcs = ["squeeze_test.cc"],

tensorflow/lite/micro/kernels/Makefile.inc

@@ -102,6 +102,7 @@ tensorflow/lite/micro/kernels/space_to_batch_nd_test.cc \
 tensorflow/lite/micro/kernels/space_to_depth_test.cc \
 tensorflow/lite/micro/kernels/split_test.cc \
 tensorflow/lite/micro/kernels/split_v_test.cc \
+tensorflow/lite/micro/kernels/squared_difference_test.cc \
 tensorflow/lite/micro/kernels/squeeze_test.cc \
 tensorflow/lite/micro/kernels/strided_slice_test.cc \
 tensorflow/lite/micro/kernels/sub_test.cc \

tensorflow/lite/micro/kernels/micro_ops.h

@@ -76,6 +76,7 @@ TfLiteRegistration Register_SHAPE();
 TfLiteRegistration Register_SLICE();
 TfLiteRegistration Register_SPACE_TO_BATCH_ND();
 TfLiteRegistration Register_SPACE_TO_DEPTH();
+TfLiteRegistration Register_SQUARED_DIFFERENCE();
 TfLiteRegistration Register_SQUEEZE();
 TfLiteRegistration Register_SUB();
 TfLiteRegistration Register_SVDF();

tensorflow/lite/micro/kernels/squared_difference.cc

+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/binary_function.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/kernels/kernel_util.h"
+#include "tensorflow/lite/micro/micro_context.h"
+#include "tensorflow/lite/micro/micro_error_reporter.h"
+
+namespace tflite {
+namespace {
+constexpr int kInputTensor1 = 0;
+constexpr int kInputTensor2 = 1;
+constexpr int kOutputTensor = 0;
+
+struct OpData {
+  bool requires_broadcast;
+  ArithmeticParams arithmetic_params;
+};
+
+template <typename T>
+T SquaredDifference(T input1, T input2) {
+  const T difference = input1 - input2;
+  return difference * difference;
+}
+
+void* SquaredDifferenceInit(TfLiteContext* context, const char* buffer,
+                            size_t length) {
+  TFLITE_DCHECK(context->AllocatePersistentBuffer != nullptr);
+  return context->AllocatePersistentBuffer(context, sizeof(OpData));
+}
+
+TfLiteStatus SquaredDifferencePrepare(TfLiteContext* context,
+                                      TfLiteNode* node) {
+  TFLITE_DCHECK(node->user_data != nullptr);
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+  data->requires_broadcast = false;
+
+  TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
+  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
+
+  MicroContext* micro_context = GetMicroContext(context);
+
+  TfLiteTensor* input1 =
+      micro_context->AllocateTempInputTensor(node, kInputTensor1);
+  TF_LITE_ENSURE(context, input1 != nullptr);
+  TfLiteTensor* input2 =
+      micro_context->AllocateTempInputTensor(node, kInputTensor2);
+  TF_LITE_ENSURE(context, input2 != nullptr);
+  TfLiteTensor* output =
+      micro_context->AllocateTempOutputTensor(node, kOutputTensor);
+  TF_LITE_ENSURE(context, output != nullptr);
+
+  TF_LITE_ENSURE_TYPES_EQ(context, input1->type, input2->type);
+  output->type = input2->type;
+
+  // Ensure the quantization parameters are equivalent.
+  if (input1->type == kTfLiteInt8) {
+    const auto& input1_quantization_params = input1->params;
+    const auto& input2_quantization_params = input2->params;
+    const auto& output_quantization_params = output->params;
+    const int32_t integer_type_min = std::numeric_limits<int8_t>::min();
+    const int32_t integer_type_max = std::numeric_limits<int8_t>::max();
+    TF_LITE_ENSURE(context,
+                   input1_quantization_params.zero_point >= integer_type_min);
+    TF_LITE_ENSURE(context,
+                   input1_quantization_params.zero_point <= integer_type_max);
+    TF_LITE_ENSURE(context,
+                   input2_quantization_params.zero_point >= integer_type_min);
+    TF_LITE_ENSURE(context,
+                   input2_quantization_params.zero_point <= integer_type_max);
+    TF_LITE_ENSURE(context,
+                   output_quantization_params.zero_point >= integer_type_min);
+    TF_LITE_ENSURE(context,
+                   output_quantization_params.zero_point <= integer_type_max);
+    data->arithmetic_params.input1_offset =
+        -input1_quantization_params.zero_point;
+    data->arithmetic_params.input2_offset =
+        -input2_quantization_params.zero_point;
+    data->arithmetic_params.output_offset =
+        output_quantization_params.zero_point;
+
+    // shift to make integer for scales.
+    // 7 is selected so that maximum shifted result 255^2 * (1 << (7 * 2 ))
+    // does not overflow signed 32-bit integer
+    data->arithmetic_params.left_shift = 7;
+    const double twice_max_input_scale =
+        2.0 * static_cast<double>(std::max(input1_quantization_params.scale,
+                                           input2_quantization_params.scale));
+    const double real_input1_multiplier =
+        static_cast<double>(input1_quantization_params.scale) /
+        twice_max_input_scale;
+    double real_input2_multiplier =
+        static_cast<double>(input2_quantization_params.scale) /
+        twice_max_input_scale;
+    const double real_output_multiplier =
+        (twice_max_input_scale * twice_max_input_scale) /
+        static_cast<double>((1 << data->arithmetic_params.left_shift * 2) *
+                            output_quantization_params.scale);
+    QuantizeMultiplierSmallerThanOneExp(
+        real_input1_multiplier, &data->arithmetic_params.input1_multiplier,
+        &data->arithmetic_params.input1_shift);
+    QuantizeMultiplierSmallerThanOneExp(
+        real_input2_multiplier, &data->arithmetic_params.input2_multiplier,
+        &data->arithmetic_params.input2_shift);
+    QuantizeMultiplierSmallerThanOneExp(
+        real_output_multiplier, &data->arithmetic_params.output_multiplier,
+        &data->arithmetic_params.output_shift);
+    data->arithmetic_params.quantized_activation_min =
+        std::numeric_limits<int8_t>::min();
+    data->arithmetic_params.quantized_activation_max =
+        std::numeric_limits<int8_t>::max();
+  }
+
+  data->requires_broadcast = !HaveSameShapes(input1, input2);
+
+  micro_context->DeallocateTempTfLiteTensor(input1);
+  micro_context->DeallocateTempTfLiteTensor(input2);
+  micro_context->DeallocateTempTfLiteTensor(output);
+  return kTfLiteOk;
+}
+
+inline int8_t SquaredDifference(int8_t x, int8_t y,
+                                const ArithmeticParams& params) {
+  const int32_t input1_val = params.input1_offset + x;
+  const int32_t input2_val = params.input2_offset + y;
+  const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
+  const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
+  const int32_t scaled_input1_val =
+      MultiplyByQuantizedMultiplierSmallerThanOneExp(
+          shifted_input1_val, params.input1_multiplier, params.input1_shift);
+  const int32_t scaled_input2_val =
+      MultiplyByQuantizedMultiplierSmallerThanOneExp(
+          shifted_input2_val, params.input2_multiplier, params.input2_shift);
+  const int32_t raw_diff = scaled_input1_val - scaled_input2_val;
+
+  // Max of this is 255^2 * (1 << 14), so won't overflow 32 bits.
+  const int32_t squared_raw_diff = raw_diff * raw_diff;
+  const int32_t raw_output =
+      MultiplyByQuantizedMultiplierSmallerThanOneExp(
+          squared_raw_diff, params.output_multiplier, params.output_shift) +
+      params.output_offset;
+  const int32_t clamped_output =
+      std::min(params.quantized_activation_max,
+               std::max(params.quantized_activation_min, raw_output));
+  return static_cast<int8_t>(clamped_output);
+}
+
+template <typename T>
+void EvalQuantizedSquaredDifference(TfLiteContext* context, TfLiteNode* node,
+                                    const OpData* data,
+                                    const TfLiteEvalTensor* input1,
+                                    const TfLiteEvalTensor* input2,
+                                    TfLiteEvalTensor* output) {
+  const auto* op_data = static_cast<const OpData*>(node->user_data);
+  if (data->requires_broadcast) {
+    reference_integer_ops::BroadcastBinaryFunction4DSlow(
+        op_data->arithmetic_params, tflite::micro::GetTensorShape(input1),
+        tflite::micro::GetTensorData<T>(input1),
+        tflite::micro::GetTensorShape(input2),
+        tflite::micro::GetTensorData<T>(input2),
+        tflite::micro::GetTensorShape(output),
+        tflite::micro::GetTensorData<T>(output),
+        reference_integer_ops::CheckArithmeticParams, SquaredDifference);
+  } else {
+    const int flat_size = tflite::micro::GetTensorShape(input1).FlatSize();
+    reference_integer_ops::ElementWise(
+        flat_size, op_data->arithmetic_params,
+        tflite::micro::GetTensorData<int8_t>(input1),
+        tflite::micro::GetTensorData<int8_t>(input2),
+        tflite::micro::GetTensorData<int8_t>(output),
+        reference_integer_ops::CheckArithmeticParams, SquaredDifference);
+  }
+}
+
+template <typename T>
+void EvalSquaredDifference(TfLiteContext* context, TfLiteNode* node,
+                           const OpData* data, const TfLiteEvalTensor* input1,
+                           const TfLiteEvalTensor* input2,
+                           TfLiteEvalTensor* output) {
+  if (data->requires_broadcast) {
+    reference_ops::BroadcastBinaryFunction4DSlow<T, T, T>(
+        tflite::micro::GetTensorShape(input1),
+        tflite::micro::GetTensorData<T>(input1),
+        tflite::micro::GetTensorShape(input2),
+        tflite::micro::GetTensorData<T>(input2),
+        tflite::micro::GetTensorShape(output),
+        tflite::micro::GetTensorData<T>(output), SquaredDifference<T>);
+  } else {
+    reference_ops::BinaryFunction<T, T, T>(
+        tflite::micro::GetTensorShape(input1),
+        tflite::micro::GetTensorData<T>(input1),
+        tflite::micro::GetTensorShape(input2),
+        tflite::micro::GetTensorData<T>(input2),
+        tflite::micro::GetTensorShape(output),
+        tflite::micro::GetTensorData<T>(output), SquaredDifference<T>);
+  }
+}
+
+TfLiteStatus SquaredDifferenceEval(TfLiteContext* context, TfLiteNode* node) {
+  OpData* data = reinterpret_cast<OpData*>(node->user_data);
+
+  const TfLiteEvalTensor* input1 =
+      tflite::micro::GetEvalInput(context, node, kInputTensor1);
+  const TfLiteEvalTensor* input2 =
+      tflite::micro::GetEvalInput(context, node, kInputTensor2);
+  TfLiteEvalTensor* output =
+      tflite::micro::GetEvalOutput(context, node, kOutputTensor);
+
+  if (output->type == kTfLiteFloat32) {
+    EvalSquaredDifference<float>(context, node, data, input1, input2, output);
+  } else if (output->type == kTfLiteInt32) {
+    EvalSquaredDifference<int32_t>(context, node, data, input1, input2, output);
+  } else if (output->type == kTfLiteInt8) {
+    EvalQuantizedSquaredDifference<int8_t>(context, node, data, input1, input2,
+                                           output);
+  } else {
+    MicroPrintf(
+        "SquaredDifference only supports FLOAT32, INT32 and INT8 now, got %d.",
+        output->type);
+    return kTfLiteError;
+  }
+
+  return kTfLiteOk;
+}
+}  // namespace
+
+TfLiteRegistration Register_SQUARED_DIFFERENCE() {
+  return tflite::micro::RegisterOp(
+      SquaredDifferenceInit, SquaredDifferencePrepare, SquaredDifferenceEval);
+}
+
+}  // namespace tflite

tensorflow/lite/micro/kernels/squared_difference_test.cc

+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/micro/kernels/kernel_runner.h"
+#include "tensorflow/lite/micro/micro_utils.h"
+#include "tensorflow/lite/micro/test_helpers.h"
+#include "tensorflow/lite/micro/testing/micro_test.h"
+
+namespace tflite {
+namespace testing {
+namespace {
+
+constexpr int kNumTestShapes = 4;
+constexpr int kMaxTestShapeSize = 5;
+
+int test_shape[kNumTestShapes][kMaxTestShapeSize] = {
+    {1, 6},
+    {2, 2, 3},
+    {3, 2, 1, 3},
+    {4, 1, 3, 1, 2},
+};
+
+template <typename T>
+void ValidateSquaredDifferenceGoldens(TfLiteTensor* tensors, int tensors_size,
+                                      const T* golden, T* output,
+                                      int output_size, float tolerance = 1e-5) {
+  int inputs_array_data[] = {2, 0, 1};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 2};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+  const TfLiteRegistration registration = tflite::Register_SQUARED_DIFFERENCE();
+  micro::KernelRunner runner(registration, tensors, tensors_size, inputs_array,
+                             outputs_array, /*builtin_data=*/nullptr);
+
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.Invoke());
+
+  for (int i = 0; i < output_size; ++i) {
+    TF_LITE_MICRO_EXPECT_NEAR(golden[i], output[i], tolerance);
+  }
+}
+
+template <typename T>
+void TestSquaredDifference(int* input1_dims_data, const T* input1_data,
+                           int* input2_dims_data, const T* input2_data,
+                           int* output_dims_data, const T* expected_output,
+                           T* output_data) {
+  TfLiteIntArray* input1_dims = IntArrayFromInts(input1_dims_data);
+  TfLiteIntArray* input2_dims = IntArrayFromInts(input2_dims_data);
+  TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+
+  constexpr int inputs_size = 2;
+  constexpr int outputs_size = 1;
+  constexpr int tensors_size = inputs_size + outputs_size;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateTensor(input1_data, input1_dims),
+      CreateTensor(input2_data, input2_dims),
+      CreateTensor(output_data, output_dims),
+  };
+
+  ValidateSquaredDifferenceGoldens(tensors, tensors_size, expected_output,
+                                   output_data, ElementCount(*output_dims));
+}
+
+template <typename T>
+void TestSquaredDifferenceQuantized(
+    int* input1_dims_data, const float* input1_data, T* input1_quantized,
+    float input1_min, float input1_max,
+
+    int* input2_dims_data, const float* input2_data, T* input2_quantized,
+    float input2_min, float input2_max,
+
+    int* output_dims_data, T* output_data, float output_min, float output_max,
+    float* dequantized_output, const float* golden,
+
+    float tolerance) {
+  QuantizationParams input1_qparams;
+  QuantizationParams input2_qparams;
+  QuantizationParams output_qparams;
+
+  input1_qparams = ChooseQuantizationParams<T>(static_cast<double>(input1_min),
+                                               static_cast<double>(input1_max));
+  input2_qparams = ChooseQuantizationParams<T>(static_cast<double>(input2_min),
+                                               static_cast<double>(input2_max));
+  output_qparams = ChooseQuantizationParams<T>(static_cast<double>(output_min),
+                                               static_cast<double>(output_max));
+
+  TfLiteIntArray* input1_dims = IntArrayFromInts(input1_dims_data);
+  TfLiteIntArray* input2_dims = IntArrayFromInts(input2_dims_data);
+  TfLiteIntArray* output_dims = IntArrayFromInts(output_dims_data);
+  int output_size = ElementCount(*output_dims);
+
+  constexpr int inputs_size = 2;
+  constexpr int outputs_size = 1;
+  constexpr int tensors_size = inputs_size + outputs_size;
+  TfLiteTensor tensors[tensors_size] = {
+      CreateQuantizedTensor<T>(input1_data, input1_quantized, input1_dims,
+                               input1_qparams.scale, input1_qparams.zero_point),
+      CreateQuantizedTensor<T>(input2_data, input2_quantized, input2_dims,
+                               input2_qparams.scale, input2_qparams.zero_point),
+      CreateQuantizedTensor<T>(output_data, output_dims, output_qparams.scale,
+                               output_qparams.zero_point),
+  };
+
+  int inputs_array_data[] = {2, 0, 1};
+  TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data);
+  int outputs_array_data[] = {1, 2};
+  TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data);
+
+  const TfLiteRegistration registration = tflite::Register_SQUARED_DIFFERENCE();
+  micro::KernelRunner runner(registration, tensors, tensors_size, inputs_array,
+                             outputs_array, /*builtin_data=*/nullptr);
+
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.InitAndPrepare());
+  TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, runner.Invoke());
+
+  Dequantize(output_data, output_size, output_qparams.scale,
+             output_qparams.zero_point, dequantized_output);
+
+  for (int i = 0; i < output_size; ++i) {
+    TF_LITE_MICRO_EXPECT_NEAR(golden[i], dequantized_output[i], tolerance);
+  }
+}
+
+}  // namespace
+}  // namespace testing
+}  // namespace tflite
+
+TF_LITE_MICRO_TESTS_BEGIN
+
+TF_LITE_MICRO_TEST(FloatSquaredDifferenceSameShape) {
+  constexpr int data_size = 4;
+  int inout_shape[] = {4, 1, 2, 2, 1};
+  const float input1_values[] = {-0.2, 0.2, -1.2, 0.8};
+  const float input2_values[] = {0.5, 0.2, -1.5, 0.5};
+  const float golden_values[] = {0.49, 0.0, 0.09, 0.09};
+  float output_data[data_size];
+  tflite::testing::TestSquaredDifference(
+      inout_shape, input1_values, inout_shape, input2_values, inout_shape,
+      golden_values, output_data);
+}
+
+TF_LITE_MICRO_TEST(FloatSquaredDifferenceVariousShapes) {
+  constexpr int data_size = 6;
+  const float input1_values[] = {-2.0, 0.2, 0.3, 0.8, 1.1, -2.0};
+  const float input2_values[] = {1.0, 0.2, 0.6, 0.4, -1.0, -0.0};
+  const float golden_values[] = {9.0, 0.0, 0.09, 0.16, 4.41, 4.0};
+  float output_data[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifference(
+        tflite::testing::test_shape[i], input1_values,
+        tflite::testing::test_shape[i], input2_values,
+        tflite::testing::test_shape[i], golden_values, output_data);
+  }
+}
+
+TF_LITE_MICRO_TEST(FloatSquaredDifferenceWithBroadcast) {
+  constexpr int data_size = 6;
+
+  // input 2 is scalar
+  int input2_shape[] = {1, 1};
+  const float input1_values[] = {-0.2, 0.2, 0.5, 0.8, 0.11, 1.1};
+  const float input2_values[] = {0.1};
+  const float golden_values[] = {0.09, 0.01, 0.16, 0.49, 0.0001, 1.0};
+  float output_data[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifference(
+        tflite::testing::test_shape[i], input1_values, input2_shape,
+        input2_values, tflite::testing::test_shape[i], golden_values,
+        output_data);
+  }
+}
+
+TF_LITE_MICRO_TEST(IntegerSquaredDifferenceSameShape) {
+  constexpr int data_size = 4;
+  int inout_shape[] = {4, 1, 2, 2, 1};
+  const int32_t input1_values[] = {-2, 2, -15, 8};
+  const int32_t input2_values[] = {5, -2, -3, 5};
+  const int32_t golden_values[] = {49, 16, 144, 9};
+  int32_t output_data[data_size];
+  tflite::testing::TestSquaredDifference(
+      inout_shape, input1_values, inout_shape, input2_values, inout_shape,
+      golden_values, output_data);
+}
+
+TF_LITE_MICRO_TEST(IntegerSquaredDifferenceVariousShapes) {
+  constexpr int data_size = 6;
+  const int32_t input1_values[] = {-20, 2, 3, 8, 11, -20};
+  const int32_t input2_values[] = {1, 2, 6, 5, -5, -20};
+  const int32_t golden_values[] = {441, 0, 9, 9, 256, 0};
+  int32_t output_data[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifference(
+        tflite::testing::test_shape[i], input1_values,
+        tflite::testing::test_shape[i], input2_values,
+        tflite::testing::test_shape[i], golden_values, output_data);
+  }
+}
+
+TF_LITE_MICRO_TEST(IntegerSquaredDifferenceWithBroadcast) {
+  constexpr int data_size = 6;
+
+  // input 2 is a scalar
+  int input2_shape[] = {1, 1};
+  const int32_t input1_values[] = {-20, 10, 7, 3, 1, 13};
+  const int32_t input2_values[] = {3};
+  const int32_t golden_values[] = {529, 49, 16, 0, 4, 100};
+  int32_t output_data[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifference(
+        tflite::testing::test_shape[i], input1_values, input2_shape,
+        input2_values, tflite::testing::test_shape[i], golden_values,
+        output_data);
+  }
+}
+
+TF_LITE_MICRO_TEST(QuantizedSquaredDifferenceSameShape) {
+  constexpr int data_size = 4;
+  int inout_shape[] = {4, 1, 2, 2, 1};
+  const float input1_values[] = {-0.2, 0.2, -1.2, 0.8};
+  const float input2_values[] = {0.5, 0.2, -1.5, 0.5};
+  const float golden_values[] = {0.49, 0.0, 0.09, 0.09};
+  int8_t input1_quantized[data_size];
+  int8_t input2_quantized[data_size];
+  int8_t output[data_size];
+  float output_dequantized[data_size];
+  tflite::testing::TestSquaredDifferenceQuantized(
+      inout_shape, input1_values, input1_quantized, -1.2f, 0.8f, inout_shape,
+      input2_values, input2_quantized, -1.5f, 0.5f, inout_shape, output, 0.0f,
+      0.5f, output_dequantized, golden_values, 2.0f / 255.0f);
+}
+
+TF_LITE_MICRO_TEST(QuantizedSquaredDifferenceVariousShapes) {
+  constexpr int data_size = 6;
+  const float input1_values[] = {-2.0, 0.2, 0.3, 0.8, 1.1, -2.0};
+  const float input2_values[] = {1.0, 0.2, 0.6, 0.4, -1.0, -0.0};
+  const float golden_values[] = {9.0, 0.0, 0.09, 0.16, 4.41, 4.0};
+  int8_t input1_quantized[data_size];
+  int8_t input2_quantized[data_size];
+  int8_t output[data_size];
+  float output_dequantized[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifferenceQuantized(
+        tflite::testing::test_shape[i], input1_values, input1_quantized, -2.0f,
+        1.7f, tflite::testing::test_shape[i], input2_values, input2_quantized,
+        -1.0f, 1.0f, tflite::testing::test_shape[i], output, 0.0f, 9.0f,
+        output_dequantized, golden_values, 18.0f / 255.0f);
+  }
+}
+
+TF_LITE_MICRO_TEST(FloatSquaredDifferenceWithBroadcast) {
+  constexpr int data_size = 6;
+
+  // input 2 is a scalar
+  int input2_shape[] = {1, 1};
+  const float input1_values[] = {-0.2, 0.2, 0.5, 0.8, 0.11, 1.1};
+  const float input2_values[] = {0.1};
+  const float golden_values[] = {0.09, 0.01, 0.16, 0.49, 0.0001, 1.0};
+  int8_t input1_quantized[data_size];
+  int8_t input2_quantized[1];
+  int8_t output[data_size];
+  float output_dequantized[data_size];
+  for (int i = 0; i < tflite::testing::kNumTestShapes; ++i) {
+    tflite::testing::TestSquaredDifferenceQuantized(
+        tflite::testing::test_shape[i], input1_values, input1_quantized, -0.2f,
+        1.1f, input2_shape, input2_values, input2_quantized, 0.0f, 1.0f,
+        tflite::testing::test_shape[i], output, 0.0f, 1.0f, output_dequantized,
+        golden_values, 2.0f / 255.0f);
+  }
+}
+
+TF_LITE_MICRO_TESTS_END

tensorflow/lite/micro/tools/make/Makefile

@@ -377,6 +377,7 @@ tensorflow/lite/micro/kernels/space_to_batch_nd.cc \
 tensorflow/lite/micro/kernels/space_to_depth.cc \
 tensorflow/lite/micro/kernels/split.cc \
 tensorflow/lite/micro/kernels/split_v.cc \
+tensorflow/lite/micro/kernels/squared_difference.cc \
 tensorflow/lite/micro/kernels/squeeze.cc \
 tensorflow/lite/micro/kernels/strided_slice.cc \
 tensorflow/lite/micro/kernels/sub.cc \