chore: import upstream snapshot with attribution

tensorflow/lite/experimental/shlo/legacy/test/BUILD

@@ -0,0 +1,207 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+cc_library(
+    name = "matchers",
+    testonly = True,
+    hdrs = ["matchers.h"],
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "util",
+    srcs = [
+    ],
+    hdrs = [
+        "util.h",
+    ],
+    tags = ["no_oss"],
+    visibility = ["//tensorflow/lite/experimental/shlo/legacy/bench:__subpackages__"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+cc_test(
+    name = "broadcast_in_dim",
+    srcs = [
+        "broadcast_in_dim_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "clamp",
+    srcs = [
+        "clamp_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "compare",
+    srcs = [
+        "compare_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "concatenate",
+    srcs = [
+        "concatenate_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "elementwise_binary",
+    srcs = [
+        "elementwise_binary_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":matchers",
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "elementwise_unary",
+    srcs = [
+        "elementwise_unary_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":matchers",
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "iota",
+    srcs = [
+        "iota_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "is_finite",
+    srcs = [
+        "is_finite_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "select",
+    srcs = [
+        "select_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "uniform_dequantize_quantize",
+    srcs = [
+        "uniform_dequantize_quantize_test.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":matchers",
+        "//tensorflow/lite/experimental/shlo/legacy:debug",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "shape_test",
+    srcs = ["shape_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/shlo/legacy/test/broadcast_in_dim_test.cc

@@ -0,0 +1,147 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(std::initializer_list<DimensionSize>&& operand_shape,
+          std::vector<typename Storage<element_type>::Type>&& operand_values,
+          std::initializer_list<DimensionSize>&& broadcast_dimensions_values,
+          std::initializer_list<DimensionSize>&& result_shape,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Tensor operand(TensorType(Shape(operand_shape), element_type),
+                 operand_values.data());
+  Tensor expected(TensorType(Shape(result_shape), element_type),
+                  expected_values.data());
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(result_shape), element_type),
+                result_values.data());
+
+  absl::Span<const DimensionSize> broadcast_dimensions(
+      broadcast_dimensions_values);
+
+  ASSERT_OK(BroadcastInDim(operand, broadcast_dimensions, result));
+  EXPECT_EQ(result, expected)
+      << "operand: " << operand
+      << "\nbroadcast_dimensions: " << ToString(broadcast_dimensions);
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    std::initializer_list<DimensionSize>&& operand_shape,
+    std::vector<typename Storage<expressed_type>::Type>&& operand_values,
+    std::initializer_list<DimensionSize>&& broadcast_dimensions_values,
+    std::initializer_list<DimensionSize>&& result_shape,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  auto operand_quant_values = QuantizeVector<storage_type, expressed_type>(
+      operand_values, quantized_parameter);
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  std::vector<typename Storage<storage_type>::Type> result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+  QuantizedTensor operand(
+      QuantizedTensorType(Shape(operand_shape),
+                          QuantizedTensorElementType(element_type)),
+      operand_quant_values.data());
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(result_shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(result_shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  absl::Span<const DimensionSize> broadcast_dimensions(
+      broadcast_dimensions_values);
+  auto res = BroadcastInDim(operand, broadcast_dimensions, result);
+
+  ASSERT_OK(BroadcastInDim(operand, broadcast_dimensions, result));
+  EXPECT_EQ(result, expected)
+      << "operand: " << operand
+      << "\nbroadcast_dimensions: " << ToString(broadcast_dimensions);
+}
+
+TEST(BroadcastInDim, Unquantized) {
+  test<ElementType::kI1>({1, 3}, {true, false, true}, {2, 1}, {2, 3, 2},
+                         {true, true, false, false, true, true, true, true,
+                          false, false, true, true});
+  test<ElementType::kSI8>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                          {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI16>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                           {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI32>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                           {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kBF16>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                           {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kF16>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                          {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kF32>({1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+                          {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+}
+
+TEST(BroadcastInDim, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {1, 3}, {1, 2, 3}, {2, 1}, {2, 3, 2},
+      {1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/clamp_test.cc

@@ -0,0 +1,173 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(std::initializer_list<DimensionSize>&& shape,
+          std::vector<typename Storage<element_type>::Type>&& min_values,
+          std::vector<typename Storage<element_type>::Type>&& operand_values,
+          std::vector<typename Storage<element_type>::Type>&& max_values,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Shape min_shape = (min_values.size() > 1) ? Shape(shape) : Shape();
+  Tensor min(TensorType(std::move(min_shape), element_type), min_values.data());
+  Shape max_shape = (max_values.size() > 1) ? Shape(shape) : Shape();
+  Tensor max(TensorType(std::move(max_shape), element_type), max_values.data());
+  Tensor operand(TensorType(Shape(shape), element_type), operand_values.data());
+  Tensor expected(TensorType(Shape(shape), element_type),
+                  expected_values.data());
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), element_type), result_values.data());
+
+  ASSERT_OK(Clamp(min, operand, max, result));
+  EXPECT_EQ(result, expected)
+      << "min: " << min << "\nmax: " << max << "\noperand: " << operand;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<expressed_type>::Type>&& min_values,
+    std::vector<typename Storage<expressed_type>::Type>&& operand_values,
+    std::vector<typename Storage<expressed_type>::Type>&& max_values,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  auto min_quant_values = QuantizeVector<storage_type, expressed_type>(
+      min_values, quantized_parameter);
+  auto operand_quant_values = QuantizeVector<storage_type, expressed_type>(
+      operand_values, quantized_parameter);
+  auto max_quant_values = QuantizeVector<storage_type, expressed_type>(
+      max_values, quantized_parameter);
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  std::vector<typename Storage<storage_type>::Type> result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  Shape min_shape = (min_values.size() > 1) ? Shape(shape) : Shape();
+  QuantizedTensor min(
+      QuantizedTensorType(std::move(min_shape),
+                          QuantizedTensorElementType(element_type)),
+      min_quant_values.data());
+  Shape max_shape = (max_values.size() > 1) ? Shape(shape) : Shape();
+  QuantizedTensor max(
+      QuantizedTensorType(std::move(max_shape),
+                          QuantizedTensorElementType(element_type)),
+      max_quant_values.data());
+  QuantizedTensor operand(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      operand_quant_values.data());
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  ASSERT_OK(Clamp(min, operand, max, result));
+  EXPECT_EQ(result, expected)
+      << "min: " << min << "\nmax: " << max << "\noperand: " << operand;
+}
+
+TEST(Clamp, Unquantized) {
+  test<ElementType::kSI8>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI16>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI32>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kBF16>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kF16>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kF32>({3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+
+  test<ElementType::kSI8>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI16>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI32>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kBF16>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kF16>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kF32>({3}, {0, 1, 1}, {-3, 0, 3}, {1, 1, 2}, {0, 1, 2});
+}
+
+TEST(Clamp, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {3}, {0}, {-2, 0, 2}, {1}, {0, 0, 1});
+
+  test<ElementType::kSI8, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                              {3}, {0, 1, 1}, {-3, 0, 3},
+                                              {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI8, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                             {3}, {0, 1, 1}, {-3, 0, 3},
+                                             {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI8, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                             {3}, {0, 1, 1}, {-3, 0, 3},
+                                             {1, 1, 2}, {0, 1, 2});
+
+  test<ElementType::kSI16, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                               {3}, {0, 1, 1}, {-3, 0, 3},
+                                               {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI16, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {3}, {0, 1, 1}, {-3, 0, 3},
+                                              {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI16, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                              {3}, {0, 1, 1}, {-3, 0, 3},
+                                              {1, 1, 2}, {0, 1, 2});
+
+  test<ElementType::kSI32, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                               {3}, {0, 1, 1}, {-3, 0, 3},
+                                               {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI32, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {3}, {0, 1, 1}, {-3, 0, 3},
+                                              {1, 1, 2}, {0, 1, 2});
+  test<ElementType::kSI32, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                              {3}, {0, 1, 1}, {-3, 0, 3},
+                                              {1, 1, 2}, {0, 1, 2});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/compare_test.cc

@@ -0,0 +1,452 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(
+    ComparisonDirection comparison_direction, CompareType compare_type,
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<element_type>::Type>&& lhs_values,
+    std::vector<typename Storage<element_type>::Type>&& rhs_values,
+    std::vector<typename Storage<ElementType::kI1>::Type>&& expected_values) {
+  Tensor lhs(TensorType(Shape(shape), element_type), lhs_values.data());
+  Tensor rhs(TensorType(Shape(shape), element_type), rhs_values.data());
+  Tensor expected(TensorType(Shape(shape), ElementType::kI1),
+                  expected_values.data());
+
+  std::vector<typename Storage<ElementType::kI1>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), ElementType::kI1),
+                result_values.data());
+
+  ASSERT_OK(Compare(lhs, rhs, comparison_direction, compare_type, result));
+  EXPECT_EQ(result, expected)
+      << "comparison_direction: " << comparison_direction
+      << "\ncompare_type: " << compare_type << "\nlhs: " << lhs
+      << "\nrhs: " << rhs;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    ComparisonDirection comparison_direction, CompareType compare_type,
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<expressed_type>::Type>&& lhs_values,
+    std::vector<typename Storage<expressed_type>::Type>&& rhs_values,
+    std::vector<typename Storage<ElementType::kI1>::Type>&& expected_values) {
+  auto lhs_quant_values = QuantizeVector<storage_type, expressed_type>(
+      lhs_values, quantized_parameter);
+  auto rhs_quant_values = QuantizeVector<storage_type, expressed_type>(
+      rhs_values, quantized_parameter);
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+  QuantizedTensor lhs(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      lhs_quant_values.data());
+  QuantizedTensor rhs(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      rhs_quant_values.data());
+
+  Tensor expected(TensorType(Shape(shape), ElementType::kI1),
+                  expected_values.data());
+
+  std::vector<typename Storage<ElementType::kI1>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), ElementType::kI1),
+                result_values.data());
+
+  ASSERT_OK(Compare(lhs, rhs, comparison_direction, compare_type, result));
+  EXPECT_EQ(result, expected)
+      << "comparison_direction: " << comparison_direction
+      << "\ncompare_type: " << compare_type << "\nlhs: " << lhs
+      << "\nrhs: " << rhs;
+}
+
+TEST(Compare, Unquantized) {
+  test<ElementType::kI1>(ComparisonDirection::kEQ, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {true, false, false, true});
+  test<ElementType::kSI8>(ComparisonDirection::kEQ, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, false, true});
+  test<ElementType::kSI16>(ComparisonDirection::kEQ, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, false, true});
+  test<ElementType::kSI32>(ComparisonDirection::kEQ, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, false, true});
+  test<ElementType::kBF16>(ComparisonDirection::kEQ, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, false, true});
+  test<ElementType::kF16>(ComparisonDirection::kEQ, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, false, true});
+  test<ElementType::kF32>(ComparisonDirection::kEQ, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, false, true});
+
+  test<ElementType::kI1>(ComparisonDirection::kNE, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {false, true, true, false});
+  test<ElementType::kSI8>(ComparisonDirection::kNE, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, true, false});
+  test<ElementType::kSI16>(ComparisonDirection::kNE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, true, false});
+  test<ElementType::kSI32>(ComparisonDirection::kNE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, true, false});
+  test<ElementType::kBF16>(ComparisonDirection::kNE, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, true, false});
+  test<ElementType::kF16>(ComparisonDirection::kNE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, true, false});
+  test<ElementType::kF32>(ComparisonDirection::kNE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, true, false});
+
+  test<ElementType::kI1>(ComparisonDirection::kGE, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {true, false, true, true});
+  test<ElementType::kSI8>(ComparisonDirection::kGE, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, true, true});
+  test<ElementType::kSI16>(ComparisonDirection::kGE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, true, true});
+  test<ElementType::kSI32>(ComparisonDirection::kGE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, true, true});
+  test<ElementType::kBF16>(ComparisonDirection::kGE, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, false, true, true});
+  test<ElementType::kF16>(ComparisonDirection::kGE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, true, true});
+  test<ElementType::kF32>(ComparisonDirection::kGE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, false, true, true});
+
+  test<ElementType::kI1>(ComparisonDirection::kGT, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {false, false, true, false});
+  test<ElementType::kSI8>(ComparisonDirection::kGT, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, false, true, false});
+  test<ElementType::kSI16>(ComparisonDirection::kGT, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, false, true, false});
+  test<ElementType::kSI32>(ComparisonDirection::kGT, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, false, true, false});
+  test<ElementType::kBF16>(ComparisonDirection::kGT, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, false, true, false});
+  test<ElementType::kF16>(ComparisonDirection::kGT, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, false, true, false});
+  test<ElementType::kF32>(ComparisonDirection::kGT, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, false, true, false});
+
+  test<ElementType::kI1>(ComparisonDirection::kLE, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {true, true, false, true});
+  test<ElementType::kSI8>(ComparisonDirection::kLE, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, true, false, true});
+  test<ElementType::kSI16>(ComparisonDirection::kLE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, true, false, true});
+  test<ElementType::kSI32>(ComparisonDirection::kLE, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, true, false, true});
+  test<ElementType::kBF16>(ComparisonDirection::kLE, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {true, true, false, true});
+  test<ElementType::kF16>(ComparisonDirection::kLE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, true, false, true});
+  test<ElementType::kF32>(ComparisonDirection::kLE, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {true, true, false, true});
+
+  test<ElementType::kI1>(ComparisonDirection::kLT, CompareType::kUnsigned, {4},
+                         {true, false, true, false}, {true, true, false, false},
+                         {false, true, false, false});
+  test<ElementType::kSI8>(ComparisonDirection::kLT, CompareType::kSigned, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, false, false});
+  test<ElementType::kSI16>(ComparisonDirection::kLT, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, false, false});
+  test<ElementType::kSI32>(ComparisonDirection::kLT, CompareType::kSigned, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, false, false});
+  test<ElementType::kBF16>(ComparisonDirection::kLT, CompareType::kFloat, {4},
+                           {1, 0, 1, 0}, {1, 1, 0, 0},
+                           {false, true, false, false});
+  test<ElementType::kF16>(ComparisonDirection::kLT, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, false, false});
+  test<ElementType::kF32>(ComparisonDirection::kLT, CompareType::kFloat, {4},
+                          {1, 0, 1, 0}, {1, 1, 0, 0},
+                          {false, true, false, false});
+}
+
+TEST(Compare, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kEQ,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, false, true});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kNE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, true, false});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, false, true, true});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kGT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, false, true, false});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLE,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {true, true, false, true});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, ComparisonDirection::kLT,
+      CompareType::kFloat, {4}, {1, 0, 1, 0}, {1, 1, 0, 0},
+      {false, true, false, false});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/concatenate_test.cc

@@ -0,0 +1,233 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstddef>
+#include <initializer_list>
+#include <sstream>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+struct TensorConst {
+  std::initializer_list<DimensionSize>&& shape;
+  std::vector<typename Storage<element_type>::Type>&& values;
+};
+
+template <typename T>
+std::string ToString(absl::string_view name,
+                     const std::vector<const T*>& tensors) {
+  std::ostringstream result;
+  for (size_t i = 0; i < tensors.size(); ++i) {
+    result << name << "[" << i << "]: " << *tensors[i] << "\n";
+  }
+  return result.str();
+}
+
+template <ElementType element_type>
+void test(std::initializer_list<TensorConst<element_type>>&& inputs_,
+          DimensionSize dimension, TensorConst<element_type>&& expected_) {
+  std::vector<Tensor> inputs_storage;
+  for (auto& x : inputs_) {
+    inputs_storage.emplace_back(
+        Tensor(TensorType(Shape(x.shape), element_type), x.values.data()));
+  }
+  std::vector<const Tensor*> inputs;
+  for (auto& x : inputs_storage) {
+    inputs.push_back(&x);
+  }
+
+  Tensor expected(TensorType(Shape(expected_.shape), element_type),
+                  expected_.values.data());
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected.num_elements());
+  Tensor result(TensorType(expected.type()), result_values.data());
+
+  ASSERT_OK(Concatenate(absl::Span<const Tensor*>(inputs), dimension, result));
+  EXPECT_EQ(result, expected)
+      << ToString("inputs", inputs) << "dimension: " << dimension;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(QuantizedParameter&& quantized_parameter,
+          std::initializer_list<TensorConst<expressed_type>>&& inputs_,
+          DimensionSize dimension, TensorConst<expressed_type>&& expected_) {
+  std::vector<std::vector<typename Storage<storage_type>::Type>> inputs_storage;
+  std::vector<QuantizedTensor> input_tensors;
+  std::vector<const QuantizedTensor*> inputs;
+  for (auto& x : inputs_) {
+    inputs_storage.emplace_back(QuantizeVector<storage_type, expressed_type>(
+        x.values, quantized_parameter));
+    QuantizedTensorElementType element_type(
+        storage_type, expressed_type, QuantizedParameter(quantized_parameter));
+    QuantizedTensorType type(Shape(x.shape), std::move(element_type));
+    input_tensors.emplace_back(
+        QuantizedTensor(std::move(type), inputs_storage.back().data()));
+  }
+  for (auto& t : input_tensors) {
+    inputs.push_back(&t);
+  }
+
+  auto quantized_expected_values = QuantizeVector<storage_type, expressed_type>(
+      expected_.values, quantized_parameter);
+  QuantizedTensor expected(
+      QuantizedTensorType(
+          Shape(expected_.shape),
+          QuantizedTensorElementType(storage_type, expressed_type,
+                                     QuantizedParameter(quantized_parameter))),
+      quantized_expected_values.data());
+
+  std::vector<typename Storage<storage_type>::Type> result_values(
+      expected.num_elements());
+  QuantizedTensor result(QuantizedTensorType(expected.type()),
+                         result_values.data());
+
+  ASSERT_OK(Concatenate(absl::Span<const QuantizedTensor*>(inputs), dimension,
+                        result));
+  EXPECT_EQ(result, expected)
+      << ToString("inputs", inputs) << "dimension: " << dimension;
+}
+
+TEST(Concatenate, Unquantized) {
+  test<ElementType::kI1>(
+      {{{3, 2}, {true, false, true, false, true, false}},
+       {{1, 2}, {false, true}}},
+      0, {{4, 2}, {true, false, true, false, true, false, false, true}});
+  test<ElementType::kSI8>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                          {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI16>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                           {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI32>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                           {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kBF16>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                           {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kF16>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                          {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kF32>({{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+                          {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+
+  test<ElementType::kI1>(
+      {{{2, 3}, {true, false, true, false, true, false}},
+       {{2, 1}, {true, false}}},
+      1, {{2, 4}, {true, false, true, true, false, true, false, false}});
+  test<ElementType::kSI8>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                          {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI16>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                           {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI32>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                           {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kBF16>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                           {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kF16>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                          {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kF32>({{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+                          {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+}
+
+TEST(Concatenate, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{3, 2}, {1, 2, 3, 4, 5, 6}}, {{1, 2}, {7, 8}}}, 0,
+      {{4, 2}, {1, 2, 3, 4, 5, 6, 7, 8}});
+
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0},
+      {{{2, 3}, {1, 2, 3, 4, 5, 6}}, {{2, 1}, {7, 8}}}, 1,
+      {{2, 4}, {1, 2, 3, 7, 4, 5, 6, 8}});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/elementwise_binary_test.cc

@@ -0,0 +1,530 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cmath>
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/matchers.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(absl::Status (*op)(const Tensor&, const Tensor&, Tensor&),
+          std::initializer_list<DimensionSize>&& shape,
+          std::vector<typename Storage<element_type>::Type>&& input1_values,
+          std::vector<typename Storage<element_type>::Type>&& input2_values,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Tensor input1(TensorType(Shape(shape), element_type),
+                std::data(input1_values));
+  Tensor input2(TensorType(Shape(shape), element_type),
+                std::data(input2_values));
+  Tensor expected(TensorType(Shape(shape), element_type),
+                  std::data(expected_values));
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), element_type), result_values.data());
+
+  ASSERT_OK(op(input1, input2, result));
+  EXPECT_THAT(result, IsAlmostSame(expected))
+      << "input1: " << input1 << "\ninput2: " << input2;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    absl::Status (*op)(const QuantizedTensor&, const QuantizedTensor&,
+                       QuantizedTensor&),
+    std::initializer_list<DimensionSize>&& shape,
+    QuantizedParameter&& quantized_parameter,
+    std::vector<typename Storage<expressed_type>::Type>&& input1_values,
+    std::vector<typename Storage<expressed_type>::Type>&& input2_values,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  auto input1_quant_values = QuantizeVector<storage_type, expressed_type>(
+      input1_values, quantized_parameter);
+  auto input2_quant_values = QuantizeVector<storage_type, expressed_type>(
+      input2_values, quantized_parameter);
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  std::vector<typename Storage<storage_type>::Type> result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  QuantizedTensor input1(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      input1_quant_values.data());
+  QuantizedTensor input2(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      input2_quant_values.data());
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  ASSERT_OK(op(input1, input2, result));
+  EXPECT_THAT(result, IsAlmostSame(expected))
+      << "input1: " << input1 << "\ninput2: " << input2;
+}
+
+TEST(ElementwiseBinary, Add) {
+  test<ElementType::kI1>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 1, 1});
+  test<ElementType::kSI8>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+  test<ElementType::kSI16>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+  test<ElementType::kSI32>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+  test<ElementType::kBF16>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+  test<ElementType::kF16>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+  test<ElementType::kF32>(Add, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 2, 1});
+}
+
+TEST(ElementwiseBinary, AddQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Add, {4}, {.scale = 1, .zero_point = 0}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Add, {4}, {.scale = 2, .zero_point = 2}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Add, {4}, {.scale = 0.5, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 30, -10});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Add, {4}, {.scale = 1, .zero_point = 0}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Add, {4}, {.scale = 2, .zero_point = 2}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Add, {4}, {.scale = 0.5, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 30, -10});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Add, {4}, {.scale = 1, .zero_point = 0}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Add, {4}, {.scale = 2, .zero_point = 2}, {10, 0, 20, 0}, {0, 0, 10, -10},
+      {10, 0, 30, -10});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Add, {4}, {.scale = 0.5, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 30, -10});
+}
+
+TEST(ElementwiseBinary, And) {
+  test<ElementType::kI1>(And, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {0, 0, 1, 0});
+  test<ElementType::kSI8>(And, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {1, 0, 4, 0});
+  test<ElementType::kSI16>(And, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {1, 0, 4, 0});
+  test<ElementType::kSI32>(And, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {1, 0, 4, 0});
+}
+
+TEST(ElementwiseBinary, Atan2) {
+  test<ElementType::kBF16>(Atan2, {4}, {3, 0, 5, 3}, {1, 1, 4, 1},
+                           {1.24904577239825442582f, 0, 0.89605538457134395617f,
+                            1.24904577239825442582f});
+  test<ElementType::kF16>(Atan2, {4}, {3, 0, 5, 3}, {1, 1, 4, 1},
+                          {1.24904577239825442582f, 0, 0.89605538457134395617f,
+                           1.24904577239825442582f});
+  test<ElementType::kF32>(Atan2, {4}, {3, 0, 5, 3}, {1, 1, 4, 1},
+                          {1.24904577239825442582f, 0, 0.89605538457134395617f,
+                           1.24904577239825442582f});
+}
+
+TEST(ElementwiseBinary, Atan2Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Atan2, {4}, {.scale = 1e-1, .zero_point = 0}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Atan2, {4}, {.scale = 1e-1, .zero_point = 2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Atan2, {4}, {.scale = 1e-1, .zero_point = -2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Atan2, {4}, {.scale = 1e-2, .zero_point = 0}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Atan2, {4}, {.scale = 1e-2, .zero_point = 2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Atan2, {4}, {.scale = 1e-3, .zero_point = -2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Atan2, {4}, {.scale = 1e-2, .zero_point = 0}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Atan2, {4}, {.scale = 1e-2, .zero_point = 2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Atan2, {4}, {.scale = 1e-3, .zero_point = -2}, {3, 0, 5, 3}, {1, 1, 4, 1},
+      {1.24904577239825442582f, 0, 0.89605538457134395617f,
+       1.24904577239825442582f});
+}
+
+TEST(ElementwiseBinary, Divide) {
+  test<ElementType::kSI8>(Divide, {4}, {2, 5, -3, -7}, {2, 2, 3, 3},
+                          {1, 2, -1, -2});
+  test<ElementType::kSI16>(Divide, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {11, 18, -8, 15});
+  test<ElementType::kSI32>(Divide, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {11, 18, -8, 15});
+  test<ElementType::kBF16>(Divide, {4}, {22, 53, -33, -77}, {2, 4, 4, -5},
+                           {11, 13.25, -8.25, 15.4});
+  test<ElementType::kF16>(Divide, {4}, {22, 53, -33, -77}, {2, 4, 4, -5},
+                          {11, 13.25, -8.25, 15.4});
+  test<ElementType::kF32>(Divide, {4}, {22, 53, -33, -77}, {2, 4, 4, -5},
+                          {11, 13.25, -8.25, 15.4});
+}
+
+TEST(ElementwiseBinary, DivideQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Divide, {4}, {.scale = 1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Divide, {4}, {.scale = 1, .zero_point = 5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Divide, {4}, {.scale = 1, .zero_point = -5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Divide, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Divide, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Divide, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {111, 133.25, -83.25, 155.4});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Divide, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Divide, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {11, 13.25, -8.25, 15.4});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Divide, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {111, 133.25, -83.25, 155.4});
+}
+
+TEST(ElementwiseBinary, Maximum) {
+  test<ElementType::kI1>(Maximum, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                         {1, 0, 1, 1});
+  test<ElementType::kSI8>(Maximum, {4}, {2, 5, -3, -7}, {2, 2, 3, 3},
+                          {2, 5, 3, 3});
+  test<ElementType::kSI16>(Maximum, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {22, 55, 4, -5});
+  test<ElementType::kSI32>(Maximum, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {22, 55, 4, -5});
+  test<ElementType::kBF16>(Maximum, {4}, {2.2, 5.3, -3.3, -7.7},
+                           {2.2, 4.4, 4.4, -5.5}, {2.2, 5.3, 4.4, -5.5});
+  test<ElementType::kF16>(Maximum, {4}, {22, 55, -33, -77},
+                          {2.5, 3.5, 4.5, -5.5}, {22, 55, 4.5, -5.5});
+  test<ElementType::kF32>(Maximum, {4}, {2.2, 5.3, -3.3, -7.7},
+                          {2.2, 4.4, 4.4, -5.5}, {2.2, 5.3, 4.4, -5.5});
+}
+
+TEST(ElementwiseBinary, MaximumQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Maximum, {4}, {.scale = 1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Maximum, {4}, {.scale = 1, .zero_point = 5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Maximum, {4}, {.scale = 1, .zero_point = -5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Maximum, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Maximum, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Maximum, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {222, 533, 4, -5});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Maximum, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Maximum, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {22, 53, 4, -5});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Maximum, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {222, 533, 4, -5});
+}
+
+TEST(ElementwiseBinary, Minimum) {
+  test<ElementType::kI1>(Minimum, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                         {0, 0, 1, 0});
+  test<ElementType::kSI8>(Minimum, {4}, {2, 5, -3, -7}, {2, 2, 3, 3},
+                          {2, 2, -3, -7});
+  test<ElementType::kSI16>(Minimum, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {2, 3, -33, -77});
+  test<ElementType::kSI32>(Minimum, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {2, 3, -33, -77});
+  test<ElementType::kBF16>(Minimum, {4}, {2.2, 5.3, -3.3, -7.7},
+                           {2.2, 4.4, 4.4, -5.5}, {2.2, 4.4, -3.3, -7.7});
+  test<ElementType::kF16>(Minimum, {4}, {22, 55, -33, -77},
+                          {2.5, 3.5, 4.5, -5.5}, {2.5, 3.5, -33, -77});
+  test<ElementType::kF32>(Minimum, {4}, {2.2, 5.3, -3.3, -7.7},
+                          {2.2, 4.4, 4.4, -5.5}, {2.2, 4.4, -3.3, -7.7});
+}
+
+TEST(ElementwiseBinary, MinimumQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Minimum, {4}, {.scale = 1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Minimum, {4}, {.scale = 1, .zero_point = 5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Minimum, {4}, {.scale = 1, .zero_point = -5}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Minimum, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Minimum, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Minimum, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {2, 4, -333, -777});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Minimum, {4}, {.scale = 5e-1, .zero_point = 0}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Minimum, {4}, {.scale = 1e-1, .zero_point = 10}, {22, 53, -33, -77},
+      {2, 4, 4, -5}, {2, 4, -33, -77});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Minimum, {4}, {.scale = 5e-2, .zero_point = -10}, {222, 533, -333, -777},
+      {2, 4, 4, -5}, {2, 4, -333, -777});
+}
+
+TEST(ElementwiseBinary, Multiply) {
+  test<ElementType::kI1>(Multiply, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                         {0, 0, 1, 0});
+  test<ElementType::kSI8>(Multiply, {4}, {2, 5, -3, -7}, {2, 2, 3, 3},
+                          {4, 10, -9, -21});
+  test<ElementType::kSI16>(Multiply, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {44, 165, -132, 385});
+  test<ElementType::kSI32>(Multiply, {4}, {22, 55, -33, -77}, {2, 3, 4, -5},
+                           {44, 165, -132, 385});
+  test<ElementType::kBF16>(Multiply, {4}, {2.2, 5.3, -3.3, -7.7},
+                           {2.2, 4.4, 4.4, -5.5}, {4.84, 23.32, -14.52, 42.35});
+  test<ElementType::kF16>(Multiply, {4}, {22, 55, -33, -77},
+                          {2.5, 3.5, 4.5, -5.5}, {55, 192.5, -148.5, 423.5});
+  test<ElementType::kF32>(Multiply, {4}, {2.2, 5.3, -3.3, -7.7},
+                          {2.2, 4.4, 4.4, -5.5}, {4.84, 23.32, -14.52, 42.35});
+}
+
+TEST(ElementwiseBinary, MultiplyQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Multiply, {4}, {.scale = 1e-1, .zero_point = 0}, {1.1, 2.2, -3.3, -4.4},
+      {0.1, 1, 0.5, 2.5}, {0.11, 2.2, -1.7, -11});
+}
+
+TEST(ElementwiseBinary, Or) {
+  test<ElementType::kI1>(Or, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 1, 1});
+  test<ElementType::kSI8>(Or, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {3, 0, 5, 1});
+  test<ElementType::kSI16>(Or, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {3, 0, 5, 1});
+  test<ElementType::kSI32>(Or, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {3, 0, 5, 1});
+}
+
+TEST(ElementwiseBinary, Power) {
+  test<ElementType::kSI8>(Power, {6}, {-2, 1, -3, 5, -3, 4}, {0, 1, 2, 3, 3, 2},
+                          {1, 1, 9, 125, -27, 16});
+  test<ElementType::kSI16>(Power, {6}, {-2, 1, -36, 5, 3, 5},
+                           {0, 1, 2, 3, 4, 5}, {1, 1, 1296, 125, 81, 3125});
+  test<ElementType::kSI32>(Power, {6}, {-2, 1, -36, 5, 3, 10},
+                           {0, 1, 2, 3, 4, 5}, {1, 1, 1296, 125, 81, 100000});
+  test<ElementType::kBF16>(Power, {6}, {-2, -0, -36, 5, 3, 1000},
+                           {2, 2, 1.1, 2, -1, 10},
+                           {4, 0, -NAN, 25, 0.3333333333333333f, 1e+30});
+  test<ElementType::kF16>(Power, {6}, {-2, -0, -36, 5, 3, 10000},
+                          {2, 2, 1.1, 2, -1, 10},
+                          {4, 0, -NAN, 25, 0.3333333333333333f, INFINITY});
+  test<ElementType::kF32>(Power, {6}, {-2, -0, -36, 5, 3, 10000},
+                          {2, 2, 1.1, 2, -1, 10},
+                          {4, 0, -NAN, 25, 0.3333333333333333f, INFINITY});
+}
+
+TEST(ElementwiseBinary, Remainder) {
+  test<ElementType::kSI8>(Remainder, {4}, {17, 18, 19, 20}, {3, 4, 5, 7},
+                          {2, 2, 4, 6});
+  test<ElementType::kSI16>(Remainder, {4}, {17, 18, 19, 20}, {3, 4, 5, 7},
+                           {2, 2, 4, 6});
+  test<ElementType::kSI32>(Remainder, {4}, {17, -17, 17, -17}, {3, 3, -3, -3},
+                           {2, -2, 2, -2});
+  test<ElementType::kBF16>(Remainder, {4}, {17, 18, 19, 20}, {3, 4, 5, 7},
+                           {2, 2, 4, 6});
+  test<ElementType::kF16>(Remainder, {4}, {17, -17, 17, -17}, {3, 3, -3, -3},
+                          {2, -2, 2, -2});
+  test<ElementType::kF32>(Remainder, {4}, {17.1, -17.1, 17.1, -17.1},
+                          {3, 3, -3, -3}, {2.1, -2.1, 2.1, -2.1});
+}
+
+TEST(ElementwiseBinary, RemainderQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 0}, {7.1, -7.1, 7.1, -7.1},
+      {3, 3, -3, -3}, {1.1, -1.1, 1.1, -1.1});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 0}, {7.1, -7.1, 7.1, -7.1},
+      {3, 3, -3, -3}, {1.1, -1.1, 1.1, -1.1});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 0}, {7.1, -7.1, 7.1, -7.1},
+      {3, 3, -3, -3}, {1.1, -1.1, 1.1, -1.1});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 4}, {17, 18, 19, 20},
+      {3, 4, 5, 7}, {2, 2, 4, 6});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 0}, {17, -17, 17, -17},
+      {3, 3, -3, -3}, {2, -2, 2, -2});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Remainder, {4}, {.scale = 1e-2, .zero_point = -10},
+      {17.1, -17.1, 17.1, -17.1}, {3, 3, -3, -3}, {2.1, -2.1, 2.1, -2.1});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 4}, {17, 18, 19, 20},
+      {3, 4, 5, 7}, {2, 2, 4, 6});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Remainder, {4}, {.scale = 1e-1, .zero_point = 0}, {17, -17, 17, -17},
+      {3, 3, -3, -3}, {2, -2, 2, -2});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Remainder, {4}, {.scale = 1e-2, .zero_point = -10},
+      {17.1, -17.1, 17.1, -17.1}, {3, 3, -3, -3}, {2.1, -2.1, 2.1, -2.1});
+}
+
+TEST(ElementwiseBinary, ShiftLeft) {
+  test<ElementType::kSI8>(ShiftLeft, {3}, {-1, 0, 1}, {1, 2, 3}, {-2, 0, 8});
+  test<ElementType::kSI16>(ShiftLeft, {3}, {-1, 0, 1}, {1, 2, 3}, {-2, 0, 8});
+  test<ElementType::kSI32>(ShiftLeft, {3}, {-1, 0, 1}, {1, 2, 3}, {-2, 0, 8});
+}
+
+TEST(ElementwiseBinary, ShiftRightArithmetic) {
+  test<ElementType::kSI8>(ShiftRightArithmetic, {3}, {-1, 0, 8}, {1, 2, 3},
+                          {-1, 0, 1});
+  test<ElementType::kSI16>(ShiftRightArithmetic, {3}, {-1, 0, 8}, {1, 2, 3},
+                           {-1, 0, 1});
+  test<ElementType::kSI32>(ShiftRightArithmetic, {3}, {-1, 0, 8}, {1, 2, 3},
+                           {-1, 0, 1});
+}
+
+TEST(ElementwiseBinary, ShiftRightLogical) {
+  test<ElementType::kSI8>(ShiftRightLogical, {3}, {-1, 0, 8}, {1, 2, 3},
+                          {0x7F, 0, 1});
+  test<ElementType::kSI16>(ShiftRightLogical, {3}, {-1, 0, 8}, {1, 2, 3},
+                           {0x7FFF, 0, 1});
+  test<ElementType::kSI32>(ShiftRightLogical, {3}, {-1, 0, 8}, {1, 2, 3},
+                           {0x7FFFFFFF, 0, 1});
+}
+
+TEST(ElementwiseBinary, Subtract) {
+  test<ElementType::kSI8>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                          {1, 0, 0, -1});
+  test<ElementType::kSI16>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                           {1, 0, 0, -1});
+  test<ElementType::kSI32>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                           {1, 0, 0, -1});
+  test<ElementType::kBF16>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                           {1, 0, 0, -1});
+  test<ElementType::kF16>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                          {1, 0, 0, -1});
+  test<ElementType::kF32>(Subtract, {4}, {1, 0, 1, 0}, {0, 0, 1, 1},
+                          {1, 0, 0, -1});
+}
+
+TEST(ElementwiseBinary, SubtractQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Subtract, {4}, {.scale = 1, .zero_point = 0}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Subtract, {4}, {.scale = 1, .zero_point = 2}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Subtract, {4}, {.scale = 1, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      Subtract, {4}, {.scale = 1e-1, .zero_point = 0}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI16, ElementType::kF16>(
+      Subtract, {4}, {.scale = 1e-1, .zero_point = 2}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Subtract, {4}, {.scale = 1e-1, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      Subtract, {4}, {.scale = 1e-3, .zero_point = 0}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI32, ElementType::kF16>(
+      Subtract, {4}, {.scale = 1e-3, .zero_point = 2}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+  test<ElementType::kSI32, ElementType::kF32>(
+      Subtract, {4}, {.scale = 1e-3, .zero_point = -10}, {10, 0, 20, 0},
+      {0, 0, 10, -10}, {10, 0, 10, 10});
+}
+
+TEST(ElementwiseBinary, Xor) {
+  test<ElementType::kI1>(Xor, {4}, {1, 0, 1, 0}, {0, 0, 1, 1}, {1, 0, 0, 1});
+  test<ElementType::kSI8>(Xor, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {2, 0, 1, 1});
+  test<ElementType::kSI16>(Xor, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {2, 0, 1, 1});
+  test<ElementType::kSI32>(Xor, {4}, {3, 0, 5, 0}, {1, 0, 4, 1}, {2, 0, 1, 1});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/elementwise_unary_test.cc

@@ -0,0 +1,578 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cmath>
+#include <cstdint>
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/matchers.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(absl::Status (*op)(const Tensor&, Tensor&),
+          std::initializer_list<DimensionSize>&& shape,
+          std::vector<typename Storage<element_type>::Type>&& input_values,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Tensor input(TensorType(Shape(shape), element_type), std::data(input_values));
+  Tensor expected(TensorType(Shape(shape), element_type),
+                  std::data(expected_values));
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), element_type), result_values.data());
+
+  ASSERT_OK(op(input, result));
+  EXPECT_THAT(result, IsAlmostSame(expected)) << "input: " << input;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    absl::Status (*op)(const QuantizedTensor&, QuantizedTensor&),
+    std::initializer_list<DimensionSize>&& shape,
+    QuantizedParameter&& quantized_parameter,
+    std::vector<typename Storage<expressed_type>::Type>&& input_values,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  auto input_quant_values = QuantizeVector<storage_type, expressed_type>(
+      input_values, quantized_parameter);
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  std::vector<typename Storage<storage_type>::Type> result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  QuantizedTensor input(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      input_quant_values.data());
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  ASSERT_OK(op(input, result));
+  EXPECT_THAT(result, IsAlmostSame(expected)) << "input: " << input;
+}
+
+TEST(ElementwiseUnary, Abs) {
+  test<ElementType::kSI8>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+  test<ElementType::kSI16>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+  test<ElementType::kSI32>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+  test<ElementType::kBF16>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+  test<ElementType::kF16>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+  test<ElementType::kF32>(Abs, {5}, {0, 1, -2, 3, -4}, {0, 1, 2, 3, 4});
+}
+
+TEST(ElementwiseBinary, AbsQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Abs, {5}, {.scale = 1, .zero_point = 0}, {0, 1, -2, 3, -4},
+      {0, 1, 2, 3, 4});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Abs, {5}, {.scale = 1e-1, .zero_point = 1}, {0, 1, -2, 3, -4},
+      {0, 1, 2, 3, 4});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Abs, {5}, {.scale = 1e-1, .zero_point = -1}, {0, 1, -2, 3, -4},
+      {0, 1, 2, 3, 4});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Abs, {5}, {.scale = 1e-3, .zero_point = -1}, {0, 1, -2, 3, -4},
+      {0, 1, 2, 3, 4});
+}
+
+TEST(ElementwiseUnary, Cbrt) {
+  test<ElementType::kBF16>(
+      Cbrt, {4}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+  test<ElementType::kF16>(
+      Cbrt, {4}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+  test<ElementType::kF32>(
+      Cbrt, {4}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+}
+
+TEST(ElementwiseUnary, CbrtQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Cbrt, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Cbrt, {4}, {.scale = 1e-1, .zero_point = -2}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Cbrt, {4}, {.scale = 1e-1, .zero_point = 4}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Cbrt, {4}, {.scale = 1e-1, .zero_point = 4}, {0, 1, -2, 3},
+      {0, 1, -1.25992104989487316476f, 1.44224957030740838232f});
+}
+
+TEST(ElementwiseUnary, Ceil) {
+  test<ElementType::kBF16>(Ceil, {4}, {0, 1.1, -2.7, 3.5}, {0, 2, -2, 4});
+  test<ElementType::kF16>(Ceil, {4}, {0, 1.1, -2.7, 3.5}, {0, 2, -2, 4});
+  test<ElementType::kF32>(Ceil, {4}, {0, 1.1, -2.7, 3.5}, {0, 2, -2, 4});
+}
+
+TEST(ElementwiseUnary, CeilQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Ceil, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1.1, -2.7, 3.5},
+      {0, 2, -2, 4});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Ceil, {4}, {.scale = 1e-1, .zero_point = 4}, {0, 1.1, -2.7, 3.5},
+      {0, 2, -2, 4});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Ceil, {4}, {.scale = 1e-1, .zero_point = -4}, {0, 1.1, -2.7, 3.5},
+      {0, 2, -2, 4});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Ceil, {4}, {.scale = 1e-2, .zero_point = -4}, {0, 1.11, -2.77, 3.55},
+      {0, 2, -2, 4});
+}
+
+TEST(ElementwiseUnary, Cosine) {
+  test<ElementType::kBF16>(Cosine, {4}, {0, 1.1, -1.1, 2.3},
+                           {1, 0.45359612142557738777f, 0.45359612142557738777f,
+                            -0.66627602127982419331f});
+  test<ElementType::kF16>(Cosine, {4}, {0, 1.1, -1.1, 2.3},
+                          {1, 0.45359612142557738777f, 0.45359612142557738777f,
+                           -0.66627602127982419331f});
+  test<ElementType::kF32>(Cosine, {4}, {0, 1.1, -1.1, 2.3},
+                          {1, 0.45359612142557738777f, 0.45359612142557738777f,
+                           -0.66627602127982419331f});
+}
+
+TEST(ElementwiseUnary, CosineQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Cosine, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1.1, -1.1, 2.3},
+      {1, 0.45359612142557738777f, 0.45359612142557738777f,
+       -0.66627602127982419331f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Cosine, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1.1, -1.1, 2.3},
+      {1, 0.45359612142557738777f, 0.45359612142557738777f,
+       -0.66627602127982419331f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Cosine, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1.1, -1.1, 2.3},
+      {1, 0.45359612142557738777f, 0.45359612142557738777f,
+       -0.66627602127982419331f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Cosine, {4}, {.scale = 1e-4, .zero_point = 0}, {0, 1.1, -1.1, 2.3},
+      {1, 0.45359612142557738777f, 0.45359612142557738777f,
+       -0.66627602127982419331f});
+}
+
+TEST(ElementwiseUnary, CountLeadingZeros) {
+  test<ElementType::kSI8>(CountLeadingZeros, {4}, {0, 1, 127, -1},
+                          {8, 7, 1, 0});
+  test<ElementType::kSI16>(CountLeadingZeros, {4}, {0, 1, 32767, -1},
+                           {16, 15, 1, 0});
+  test<ElementType::kSI32>(CountLeadingZeros, {4}, {0, 1, 2147483647, -1},
+                           {32, 31, 1, 0});
+}
+
+TEST(ElementwiseUnary, Exponential) {
+  test<ElementType::kBF16>(Exponential, {4}, {0, 0.5, 1, 1.5},
+                           {1, 1.64872127070012814684f, 2.71828182845904523536f,
+                            4.48168907033806482260f});
+  test<ElementType::kF16>(Exponential, {4}, {0, 0.5, 1, 1.5},
+                          {1, 1.64872127070012814684f, 2.71828182845904523536f,
+                           4.48168907033806482260f});
+  test<ElementType::kF32>(Exponential, {4}, {0, 0.5, 1, 1.5},
+                          {1, 1.64872127070012814684f, 2.71828182845904523536f,
+                           4.48168907033806482260f});
+}
+
+TEST(ElementwiseUnary, ExponentialQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Exponential, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 0.5, 1, 1.5},
+      {1, 1.64872127070012814684f, 2.71828182845904523536f,
+       4.48168907033806482260f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Exponential, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 0.5, 1, 1.5},
+      {1, 1.64872127070012814684f, 2.71828182845904523536f,
+       4.48168907033806482260f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Exponential, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 0.5, 1, 1.5},
+      {1, 1.64872127070012814684f, 2.71828182845904523536f,
+       4.48168907033806482260f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Exponential, {4}, {.scale = 1e-2, .zero_point = 0}, {0, 0.5, 1, 1.5},
+      {1, 1.64872127070012814684f, 2.71828182845904523536f,
+       4.48168907033806482260f});
+}
+
+TEST(ElementwiseUnary, ExponentialMinusOne) {
+  test<ElementType::kBF16>(ExponentialMinusOne, {4}, {0, 0.5, 1, 1.5},
+                           {0, 0.64872127070012814684f, 1.71828182845904523536f,
+                            3.48168907033806482260f});
+  test<ElementType::kF16>(ExponentialMinusOne, {4}, {0, 0.5, 1, 1.5},
+                          {0, 0.64872127070012814684f, 1.71828182845904523536f,
+                           3.48168907033806482260f});
+  test<ElementType::kF32>(ExponentialMinusOne, {4}, {0, 0.5, 1, 1.5},
+                          {0, 0.64872127070012814684f, 1.71828182845904523536f,
+                           3.48168907033806482260f});
+}
+
+TEST(ElementwiseUnary, ExponentialMinusOneQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      ExponentialMinusOne, {4}, {.scale = 1e-1, .zero_point = 0},
+      {0, 0.5, 1, 1.5},
+      {0, 0.64872127070012814684f, 1.71828182845904523536f,
+       3.48168907033806482260f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      ExponentialMinusOne, {4}, {.scale = 1e-1, .zero_point = 0},
+      {0, 0.5, 1, 1.5},
+      {0, 0.64872127070012814684f, 1.71828182845904523536f,
+       3.48168907033806482260f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      ExponentialMinusOne, {4}, {.scale = 1e-1, .zero_point = 0},
+      {0, 0.5, 1, 1.5},
+      {0, 0.64872127070012814684f, 1.71828182845904523536f,
+       3.48168907033806482260f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      ExponentialMinusOne, {4}, {.scale = 1e-2, .zero_point = 0},
+      {0, 0.5, 1, 1.5},
+      {0, 0.64872127070012814684f, 1.71828182845904523536f,
+       3.48168907033806482260f});
+}
+
+TEST(ElementwiseUnary, Floor) {
+  test<ElementType::kBF16>(Floor, {4}, {0, 1.1, -2.7, 3.5}, {0, 1, -3, 3});
+  test<ElementType::kF16>(Floor, {4}, {0, 1.1, -2.7, 3.5}, {0, 1, -3, 3});
+  test<ElementType::kF32>(Floor, {4}, {0, 1.1, -2.7, 3.5}, {0, 1, -3, 3});
+}
+
+TEST(ElementwiseUnary, FloorQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Floor, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1.1, -2.7, 3.5},
+      {0, 1, -3, 3});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Floor, {4}, {.scale = 1e-1, .zero_point = 4}, {0, 1.1, -2.7, 3.5},
+      {0, 1, -3, 3});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Floor, {4}, {.scale = 1e-1, .zero_point = -4}, {0, 1.1, -2.7, 3.5},
+      {0, 1, -3, 3});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Floor, {4}, {.scale = 1e-2, .zero_point = -4}, {0, 1.11, -2.77, 3.55},
+      {0, 1, -3, 3});
+}
+
+TEST(ElementwiseUnary, Log) {
+  test<ElementType::kBF16>(Log, {4}, {0.1, 0.5, 1, 1.5},
+                           {-2.30258509299404568401f, -0.69314718055994530941f,
+                            0, 0.40546510810816438197f});
+  test<ElementType::kF16>(Log, {4}, {0.1, 0.5, 1, 1.5},
+                          {-2.30258509299404568401f, -0.69314718055994530941f,
+                           0, 0.40546510810816438197f});
+  test<ElementType::kF32>(Log, {4}, {0.1, 0.5, 1, 1.5},
+                          {-2.30258509299404568401f, -0.69314718055994530941f,
+                           0, 0.40546510810816438197f});
+}
+
+TEST(ElementwiseUnary, LogQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Log, {4}, {.scale = 1e-1, .zero_point = -4}, {0.1, 0.5, 1, 1.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Log, {4}, {.scale = 1e-1, .zero_point = -4}, {0.1, 0.5, 1, 1.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Log, {4}, {.scale = 1e-1, .zero_point = -4}, {0.1, 0.5, 1, 1.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Log, {4}, {.scale = 1e-3, .zero_point = -4}, {0.1, 0.5, 1, 1.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+}
+
+TEST(ElementwiseUnary, LogPlusOne) {
+  test<ElementType::kBF16>(LogPlusOne, {4}, {-0.9, -0.5, 0, 0.5},
+                           {-2.30258509299404568401f, -0.69314718055994530941f,
+                            0, 0.40546510810816438197f});
+  test<ElementType::kF16>(LogPlusOne, {4}, {-0.9, -0.5, 0, 0.5},
+                          {-2.30258509299404568401f, -0.69314718055994530941f,
+                           0, 0.40546510810816438197f});
+  test<ElementType::kF32>(LogPlusOne, {4}, {-0.9, -0.5, 0, 0.5},
+                          {-2.30258509299404568401f, -0.69314718055994530941f,
+                           0, 0.40546510810816438197f});
+}
+
+TEST(ElementwiseUnary, LogPlusOneQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      LogPlusOne, {4}, {.scale = 1e-1, .zero_point = 0}, {-0.9, -0.5, 0, 0.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      LogPlusOne, {4}, {.scale = 1e-1, .zero_point = 0}, {-0.9, -0.5, 0, 0.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      LogPlusOne, {4}, {.scale = 1e-1, .zero_point = 0}, {-0.9, -0.5, 0, 0.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      LogPlusOne, {4}, {.scale = 1e-4, .zero_point = 0}, {-0.9, -0.5, 0, 0.5},
+      {-2.30258509299404568401f, -0.69314718055994530941f, 0,
+       0.40546510810816438197f});
+}
+
+TEST(ElementwiseUnary, Logistic) {
+  test<ElementType::kBF16>(Logistic, {4}, {-1, -0.5, 0, 0.5},
+                           {0.26894142136999512074f, 0.37754066879814543536f,
+                            0.5, 0.62245933120185456464f});
+  test<ElementType::kF16>(Logistic, {4}, {-1, -0.5, 0, 0.5},
+                          {0.26894142136999512074f, 0.37754066879814543536f,
+                           0.5, 0.62245933120185456464f});
+  test<ElementType::kF32>(Logistic, {4}, {-1, -0.5, 0, 0.5},
+                          {0.26894142136999512074f, 0.37754066879814543536f,
+                           0.5, 0.62245933120185456464f});
+}
+
+TEST(ElementwiseUnary, LogisticQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Logistic, {4}, {.scale = 1e-1, .zero_point = 0}, {-1, -0.5, 0, 0.5},
+      {0.26894142136999512074f, 0.37754066879814543536f, 0.5,
+       0.62245933120185456464f});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Logistic, {4}, {.scale = 1e-1, .zero_point = 0}, {-1, -0.5, 0, 0.5},
+      {0.26894142136999512074f, 0.37754066879814543536f, 0.5,
+       0.62245933120185456464f});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Logistic, {4}, {.scale = 1e-1, .zero_point = 0}, {-1, -0.5, 0, 0.5},
+      {0.26894142136999512074f, 0.37754066879814543536f, 0.5,
+       0.62245933120185456464f});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Logistic, {4}, {.scale = 1e-3, .zero_point = 0}, {-1, -0.5, 0, 0.5},
+      {0.26894142136999512074f, 0.37754066879814543536f, 0.5,
+       0.62245933120185456464f});
+}
+
+TEST(ElementwiseUnary, Negate) {
+  test<ElementType::kSI8>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+  test<ElementType::kSI16>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+  test<ElementType::kSI32>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+  test<ElementType::kBF16>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+  test<ElementType::kF16>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+  test<ElementType::kF32>(Negate, {5}, {0, 1, -2, 3, -4}, {0, -1, 2, -3, 4});
+}
+
+TEST(ElementwiseBinary, NegateQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Negate, {5}, {.scale = 1, .zero_point = 0}, {0, 1, -2, 3, -4},
+      {0, -1, 2, -3, 4});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Negate, {5}, {.scale = 1e-1, .zero_point = 1}, {0, 1, -2, 3, -4},
+      {0, -1, 2, -3, 4});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Negate, {5}, {.scale = 1e-1, .zero_point = -1}, {0, 1, -2, 3, -4},
+      {0, -1, 2, -3, 4});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Negate, {5}, {.scale = 1e-3, .zero_point = -1}, {0, 1, -2, 3, -4},
+      {0, -1, 2, -3, 4});
+}
+
+TEST(ElementwiseUnary, Not) {
+  test<ElementType::kI1>(Not, {2}, {0, 1}, {1, 0});
+  test<ElementType::kSI8>(Not, {5}, {-2, -1, 0, 1, 2},
+                          {1, 0, int8_t(0xFF), int8_t(0xFE), int8_t(0xFD)});
+  test<ElementType::kSI16>(
+      Not, {5}, {-2, -1, 0, 1, 2},
+      {1, 0, int16_t(0xFFFF), int16_t(0xFFFE), int16_t(0xFFFD)});
+  test<ElementType::kSI32>(
+      Not, {5}, {-2, -1, 0, 1, 2},
+      {1, 0, int32_t(0xFFFFFFFFU), int32_t(0xFFFFFFFEU), int32_t(0xFFFFFFFDU)});
+}
+
+TEST(ElementwiseUnary, Popcnt) {
+  test<ElementType::kSI8>(Popcnt, {4}, {0, 1, 2, 127}, {0, 1, 1, 7});
+  test<ElementType::kSI16>(Popcnt, {4}, {0, 1, 2, 127}, {0, 1, 1, 7});
+  test<ElementType::kSI32>(Popcnt, {4}, {0, 1, 2, 127}, {0, 1, 1, 7});
+}
+
+TEST(ElementwiseUnary, RoundNearestAfz) {
+  test<ElementType::kBF16>(RoundNearestAfz, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                           {-3.0, 0.0, 1.0, 1.0, 3.0});
+  test<ElementType::kF16>(RoundNearestAfz, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                          {-3.0, 0.0, 1.0, 1.0, 3.0});
+  test<ElementType::kF32>(RoundNearestAfz, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                          {-3.0, 0.0, 1.0, 1.0, 3.0});
+}
+
+TEST(ElementwiseBinary, RoundNearestAfzQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      RoundNearestAfz, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-3.0, 0.0, 1.0, 1.0, 3.0});
+  test<ElementType::kSI8, ElementType::kF16>(
+      RoundNearestAfz, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-3.0, 0.0, 1.0, 1.0, 3.0});
+  test<ElementType::kSI8, ElementType::kF32>(
+      RoundNearestAfz, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-3.0, 0.0, 1.0, 1.0, 3.0});
+  test<ElementType::kSI16, ElementType::kF32>(
+      RoundNearestAfz, {5}, {.scale = 1e-2, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-3.0, 0.0, 1.0, 1.0, 3.0});
+}
+
+TEST(ElementwiseUnary, RoundNearestEven) {
+  test<ElementType::kBF16>(RoundNearestEven, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                           {-2.0, 0.0, 0.0, 1.0, 2.0});
+  test<ElementType::kF16>(RoundNearestEven, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                          {-2.0, 0.0, 0.0, 1.0, 2.0});
+  test<ElementType::kF32>(RoundNearestEven, {5}, {-2.5, 0.4, 0.5, 0.6, 2.5},
+                          {-2.0, 0.0, 0.0, 1.0, 2.0});
+}
+
+TEST(ElementwiseBinary, RoundNearestEvenQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      RoundNearestEven, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-2.0, 0.0, 0.0, 1.0, 2.0});
+  test<ElementType::kSI8, ElementType::kF16>(
+      RoundNearestEven, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-2.0, 0.0, 0.0, 1.0, 2.0});
+  test<ElementType::kSI8, ElementType::kF32>(
+      RoundNearestEven, {5}, {.scale = 1e-1, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-2.0, 0.0, 0.0, 1.0, 2.0});
+  test<ElementType::kSI16, ElementType::kF32>(
+      RoundNearestEven, {5}, {.scale = 1e-2, .zero_point = 0},
+      {-2.5, 0.4, 0.5, 0.6, 2.5}, {-2.0, 0.0, 0.0, 1.0, 2.0});
+}
+
+TEST(ElementwiseUnary, Rsqrt) {
+  test<ElementType::kBF16>(Rsqrt, {4}, {1.0, 4.0, 9.0, 25.0},
+                           {1.0, 1.0 / 2.0, 1.0 / 3.0, 1.0 / 5.0});
+  test<ElementType::kF16>(Rsqrt, {4}, {1.0, 4.0, 9.0, 25.0},
+                          {1.0, 1.0 / 2.0, 1.0 / 3.0, 1.0 / 5.0});
+  test<ElementType::kF32>(Rsqrt, {4}, {1.0, 4.0, 9.0, 25.0},
+                          {1.0, 1.0 / 2.0, 1.0 / 3.0, 1.0 / 5.0});
+}
+
+TEST(ElementwiseUnary, RsqrtQuantized) {
+  test<ElementType::kSI16, ElementType::kF32>(
+      Rsqrt, {4}, {.scale = 1e-3, .zero_point = 0}, {1.0, 4.0, 9.0, 25.0},
+      {1.0, 1.0 / 2.0, 1.0 / 3.0, 1.0 / 5.0});
+}
+
+TEST(ElementwiseUnary, Sign) {
+  test<ElementType::kSI8>(Sign, {3}, {-2, 0, 2}, {-1, 0, 1});
+  test<ElementType::kSI16>(Sign, {3}, {-2, 0, 2}, {-1, 0, 1});
+  test<ElementType::kSI32>(Sign, {3}, {-2, 0, 2}, {-1, 0, 1});
+  test<ElementType::kBF16>(
+      Sign, {8}, {+NAN, -NAN, +INFINITY, -INFINITY, -2.0, -0.0, +0.0, 2.0},
+      {NAN, NAN, 1, -1, -1, 0, 0, 1});
+  test<ElementType::kF16>(
+      Sign, {8}, {+NAN, -NAN, +INFINITY, -INFINITY, -2.0, -0.0, +0.0, 2.0},
+      {NAN, NAN, 1, -1, -1, 0, 0, 1});
+  test<ElementType::kF32>(
+      Sign, {8}, {+NAN, -NAN, +INFINITY, -INFINITY, -2.0, -0.0, +0.0, 2.0},
+      {NAN, NAN, 1, -1, -1, 0, 0, 1});
+}
+
+TEST(ElementwiseUnary, SignQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Sign, {4}, {.scale = 1e-1, .zero_point = 0}, {-2.0, -0.0, +0.0, 2.0},
+      {-1, 0, 0, 1});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Sign, {4}, {.scale = 1e-1, .zero_point = 0}, {-2.0, -0.0, +0.0, 2.0},
+      {-1, 0, 0, 1});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Sign, {4}, {.scale = 1e-1, .zero_point = 0}, {-2.0, -0.0, +0.0, 2.0},
+      {-1, 0, 0, 1});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Sign, {4}, {.scale = 1e-2, .zero_point = 0}, {-2.0, -0.0, +0.0, 2.0},
+      {-1, 0, 0, 1});
+}
+
+TEST(ElementwiseUnary, Sine) {
+  test<ElementType::kBF16>(Sine, {5}, {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI},
+                           {0, 1, 0, -1, 0});
+  test<ElementType::kF16>(Sine, {5}, {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI},
+                          {0, 1, 0, -1, 0});
+  test<ElementType::kF32>(Sine, {5}, {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI},
+                          {0, 1, 0, -1, 0});
+}
+
+TEST(ElementwiseUnary, SineQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Sine, {5}, {.scale = 1e-1, .zero_point = 0},
+      {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI}, {0, 1, 0, -1, 0});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Sine, {5}, {.scale = 1e-1, .zero_point = 0},
+      {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI}, {0, 1, 0, -1, 0});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Sine, {5}, {.scale = 1e-1, .zero_point = 0},
+      {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI}, {0, 1, 0, -1, 0});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Sine, {5}, {.scale = 1e-2, .zero_point = 0},
+      {0, M_PI_2, M_PI, 3 * M_PI_2, 2 * M_PI}, {0, 1, 0, -1, 0});
+}
+
+TEST(ElementwiseUnary, Sqrt) {
+  test<ElementType::kBF16>(Sqrt, {4}, {0, 1, 4, 9}, {0, 1, 2, 3});
+  test<ElementType::kF16>(Sqrt, {4}, {0, 1, 4, 9}, {0, 1, 2, 3});
+  test<ElementType::kF32>(Sqrt, {4}, {0, 1, 4, 9}, {0, 1, 2, 3});
+}
+
+TEST(ElementwiseUnary, SqrtQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Sqrt, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1, 4, 9}, {0, 1, 2, 3});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Sqrt, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1, 4, 9}, {0, 1, 2, 3});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Sqrt, {4}, {.scale = 1e-1, .zero_point = 0}, {0, 1, 4, 9}, {0, 1, 2, 3});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Sqrt, {4}, {.scale = 1e-2, .zero_point = 0}, {0, 1, 4, 9}, {0, 1, 2, 3});
+}
+
+TEST(ElementwiseUnary, Tanh) {
+  test<ElementType::kBF16>(Tanh, {3}, {-1, 0, 1},
+                           {-0.76159416, 0.0, 0.76159416});
+  test<ElementType::kF16>(Tanh, {3}, {-1, 0, 1},
+                          {-0.76159416, 0.0, 0.76159416});
+  test<ElementType::kF32>(Tanh, {3}, {-1, 0, 1},
+                          {-0.76159416, 0.0, 0.76159416});
+}
+
+TEST(ElementwiseUnary, TanhQuantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      Tanh, {3}, {.scale = 1e-1, .zero_point = 0}, {-1, 0, 1},
+      {-0.76159416, 0.0, 0.76159416});
+  test<ElementType::kSI8, ElementType::kF16>(
+      Tanh, {3}, {.scale = 1e-1, .zero_point = 0}, {-1, 0, 1},
+      {-0.76159416, 0.0, 0.76159416});
+  test<ElementType::kSI8, ElementType::kF32>(
+      Tanh, {3}, {.scale = 1e-1, .zero_point = 0}, {-1, 0, 1},
+      {-0.76159416, 0.0, 0.76159416});
+  test<ElementType::kSI16, ElementType::kF32>(
+      Tanh, {3}, {.scale = 1e-2, .zero_point = 0}, {-1, 0, 1},
+      {-0.76159416, 0.0, 0.76159416});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/iota_test.cc

@@ -0,0 +1,162 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(std::initializer_list<DimensionSize>&& shape,
+          DimensionSize iota_dimension,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Tensor expected(TensorType(Shape(shape), element_type),
+                  expected_values.data());
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), element_type), result_values.data());
+
+  ASSERT_OK(Iota(iota_dimension, result));
+  EXPECT_EQ(result, expected) << "\niota_dimension: " << iota_dimension;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    std::initializer_list<DimensionSize>&& shape, DimensionSize iota_dimension,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  decltype(expected_quant_values) result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  ASSERT_OK(Iota(iota_dimension, result));
+  EXPECT_EQ(result, expected) << "\niota_dimension: " << iota_dimension;
+}
+
+TEST(Iota, Unquantized) {
+  test<ElementType::kSI8>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI8>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kSI16>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI16>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kSI32>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI32>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kBF16>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kBF16>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kF16>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kF16>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kF32>(
+      {4, 5}, 0, {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kF32>(
+      {4, 5}, 1, {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+}
+
+TEST(Iota, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 1e-3, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 1e-3, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 1e-2, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 1e-3, .zero_point = 0}, {4, 5}, 0,
+      {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 1e-3, .zero_point = 0}, {4, 5}, 1,
+      {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/is_finite_test.cc

@@ -0,0 +1,117 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cmath>
+#include <initializer_list>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<element_type>::Type>&& input_values,
+    std::vector<typename Storage<ElementType::kI1>::Type>&& expected_values) {
+  Tensor input(TensorType(Shape(shape), element_type), input_values.data());
+  Tensor expected(TensorType(Shape(shape), ElementType::kI1),
+                  expected_values.data());
+
+  std::vector<typename Storage<ElementType::kI1>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), ElementType::kI1),
+                result_values.data());
+
+  ASSERT_OK(IsFinite(input, result));
+  EXPECT_EQ(result, expected) << "input: " << input;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<expressed_type>::Type>&& input_values,
+    std::vector<typename Storage<ElementType::kI1>::Type>&& expected_values) {
+  Tensor input(TensorType(Shape(shape), expressed_type), input_values.data());
+  Tensor expected(TensorType(Shape(shape), ElementType::kI1),
+                  expected_values.data());
+
+  std::vector<typename Storage<ElementType::kI1>::Type> result_values(
+      input_values.size());
+  Tensor result(TensorType(Shape(shape), ElementType::kI1),
+                result_values.data());
+
+  ASSERT_OK(IsFinite(input, result));
+  EXPECT_EQ(result, expected) << "input: " << input;
+}
+
+TEST(IsFinite, Unquantized) {
+  test<ElementType::kBF16>({7}, {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+                           {false, false, false, false, true, true, true});
+  test<ElementType::kF16>({7}, {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+                          {false, false, false, false, true, true, true});
+  test<ElementType::kF32>({7}, {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+                          {false, false, false, false, true, true, true});
+}
+
+TEST(IsFinite, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {7},
+      {+NAN, -NAN, -INFINITY, +INFINITY, -1, 0, 1},
+      {false, false, false, false, true, true, true});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/matchers.h

@@ -0,0 +1,30 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_MATCHERS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_MATCHERS_H_
+
+#include <gmock/gmock.h>
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"
+
+namespace stablehlo {
+namespace testing {
+
+MATCHER_P(IsAlmostSame, expected, "") { return AlmostSame(arg, expected); }
+
+}  // namespace testing
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_MATCHERS_H_

tensorflow/lite/experimental/shlo/legacy/test/select_test.cc

@@ -0,0 +1,173 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/util.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType element_type>
+void test(std::initializer_list<DimensionSize>&& shape,
+          std::vector<typename Storage<ElementType::kI1>::Type>&& pred_values,
+          std::vector<typename Storage<element_type>::Type>&& on_true_values,
+          std::vector<typename Storage<element_type>::Type>&& on_false_values,
+          std::vector<typename Storage<element_type>::Type>&& expected_values) {
+  Shape pred_shape = (pred_values.size() > 1) ? Shape(shape) : Shape();
+  Tensor pred(TensorType(std::move(pred_shape), ElementType::kI1),
+              pred_values.data());
+  Tensor on_true(TensorType(Shape(shape), element_type), on_true_values.data());
+  Tensor on_false(TensorType(Shape(shape), element_type),
+                  on_false_values.data());
+  Tensor expected(TensorType(Shape(shape), element_type),
+                  expected_values.data());
+
+  std::vector<typename Storage<element_type>::Type> result_values(
+      expected_values.size());
+  Tensor result(TensorType(Shape(shape), element_type), result_values.data());
+
+  ASSERT_OK(Select(pred, on_true, on_false, result));
+  EXPECT_EQ(result, expected) << "pred: " << pred << "\non_true: " << on_true
+                              << "\nnon_false: " << on_false;
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(
+    QuantizedParameter&& quantized_parameter,
+    std::initializer_list<DimensionSize>&& shape,
+    std::vector<typename Storage<ElementType::kI1>::Type>&& pred_values,
+    std::vector<typename Storage<expressed_type>::Type>&& on_true_values,
+    std::vector<typename Storage<expressed_type>::Type>&& on_false_values,
+    std::vector<typename Storage<expressed_type>::Type>&& expected_values) {
+  Shape pred_shape = (pred_values.size() > 1) ? Shape(shape) : Shape();
+  Tensor pred(TensorType(std::move(pred_shape), ElementType::kI1),
+              pred_values.data());
+
+  auto on_true_quant_values = QuantizeVector<storage_type, expressed_type>(
+      on_true_values, quantized_parameter);
+  auto on_false_quant_values = QuantizeVector<storage_type, expressed_type>(
+      on_false_values, quantized_parameter);
+  auto expected_quant_values = QuantizeVector<storage_type, expressed_type>(
+      expected_values, quantized_parameter);
+  std::vector<typename Storage<storage_type>::Type> result_quant_values(
+      expected_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+  QuantizedTensor on_true(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      on_true_quant_values.data());
+  QuantizedTensor on_false(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      on_false_quant_values.data());
+  QuantizedTensor expected(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      expected_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  ASSERT_OK(Select(pred, on_true, on_false, result));
+  EXPECT_EQ(result, expected) << "pred: " << pred << "\non_true: " << on_true
+                              << "\nnon_false: " << on_false;
+}
+
+TEST(Select, Unquantized) {
+  test<ElementType::kI1>({2}, {true}, {true, false}, {false, true},
+                         {true, false});
+  test<ElementType::kSI8>({2}, {false}, {1, 2}, {-1, -2}, {-1, -2});
+  test<ElementType::kSI16>({2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI32>({2}, {false}, {1, 2}, {-1, -2}, {-1, -2});
+  test<ElementType::kBF16>({2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kF16>({2}, {false}, {1, 2}, {-1, -2}, {-1, -2});
+  test<ElementType::kF32>({2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+
+  test<ElementType::kI1>({2}, {true, false}, {true, true}, {false, false},
+                         {true, false});
+  test<ElementType::kSI8>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+  test<ElementType::kSI16>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+  test<ElementType::kSI32>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+  test<ElementType::kBF16>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+  test<ElementType::kF16>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+  test<ElementType::kF32>({2}, {true, false}, {1, 2}, {-1, -2}, {1, -2});
+}
+
+TEST(Select, Quantized) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI8, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                             {2}, {false}, {1, 2}, {-1, -2},
+                                             {-1, -2});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI8, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {true, false}, {1, 2},
+                                              {-1, -2}, {1, -2});
+  test<ElementType::kSI8, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                             {2}, {true, false}, {1, 2},
+                                             {-1, -2}, {1, -2});
+  test<ElementType::kSI8, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                             {2}, {true, false}, {1, 2},
+                                             {-1, -2}, {1, -2});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI16, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {false}, {1, 2}, {-1, -2},
+                                              {-1, -2});
+  test<ElementType::kSI16, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI16, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                               {2}, {true, false}, {1, 2},
+                                               {-1, -2}, {1, -2});
+  test<ElementType::kSI16, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {true, false}, {1, 2},
+                                              {-1, -2}, {1, -2});
+  test<ElementType::kSI16, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {true, false}, {1, 2},
+                                              {-1, -2}, {1, -2});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI32, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {false}, {1, 2}, {-1, -2},
+                                              {-1, -2});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {.scale = 0.1, .zero_point = 0}, {2}, {true}, {1, 2}, {-1, -2}, {1, 2});
+  test<ElementType::kSI32, ElementType::kBF16>({.scale = 0.1, .zero_point = 0},
+                                               {2}, {true, false}, {1, 2},
+                                               {-1, -2}, {1, -2});
+  test<ElementType::kSI32, ElementType::kF16>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {true, false}, {1, 2},
+                                              {-1, -2}, {1, -2});
+  test<ElementType::kSI32, ElementType::kF32>({.scale = 0.1, .zero_point = 0},
+                                              {2}, {true, false}, {1, 2},
+                                              {-1, -2}, {1, -2});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/shape_test.cc

@@ -0,0 +1,38 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+
+namespace stablehlo {
+namespace {
+
+TEST(ShapeTest, NumElementsEmpty) {
+  Shape shape;
+  EXPECT_EQ(shape.num_elements(), 0);
+}
+
+TEST(ShapeTest, NumElementsNonEmpty) {
+  Shape shape({2, 3});
+  EXPECT_EQ(shape.num_elements(), 6);
+}
+
+TEST(ShapeTest, NumElementsLargeDimensions) {
+  Shape shape({65536ULL, 45876ULL});
+  EXPECT_EQ(shape.num_elements(), 3006529536ULL);
+}
+
+}  // namespace
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/uniform_dequantize_quantize_test.cc

@@ -0,0 +1,111 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <initializer_list>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"  // IWYU pragma: keep, b/321245930
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/test/matchers.h"
+
+namespace stablehlo {
+namespace testing {
+
+template <ElementType storage_type, ElementType expressed_type>
+void test(std::initializer_list<DimensionSize>&& shape,
+          QuantizedParameter&& quantized_parameter,
+          std::vector<typename Storage<expressed_type>::Type>&& input_values) {
+  Tensor input(TensorType(Shape(shape), expressed_type), input_values.data());
+
+  std::vector<typename Storage<storage_type>::Type> quant_values(
+      input_values.size());
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+  QuantizedTensor quant(
+      QuantizedTensorType(Shape(shape), std::move(element_type)),
+      quant_values.data());
+
+  std::vector<typename Storage<expressed_type>::Type> result_values(
+      input_values.size());
+  Tensor result(TensorType(Shape(shape), expressed_type), result_values.data());
+
+  ASSERT_OK(UniformQuantize(input, quant));
+  ASSERT_OK(UniformDequantize(quant, result));
+  EXPECT_THAT(result, IsAlmostSame(input));
+}
+
+TEST(QuantizeDequantize, All) {
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {4}, {.scale = 1, .zero_point = 0}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {4}, {.scale = 1, .zero_point = 0}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI8, ElementType::kBF16>(
+      {4}, {.scale = 1e-1, .zero_point = -5}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI8, ElementType::kF16>({4}, {.scale = 1, .zero_point = 5},
+                                             {-2, -1, 0, 1, 2});
+  test<ElementType::kSI8, ElementType::kF16>(
+      {4}, {.scale = 1e-1, .zero_point = -10}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI8, ElementType::kF32>({4}, {.scale = 1, .zero_point = 5},
+                                             {-2, -1, 0, 1, 2});
+  test<ElementType::kSI8, ElementType::kF32>(
+      {4}, {.scale = 1e-1, .zero_point = +10}, {-2.2, -1.1, 0, 1.1, 2.2});
+
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {4}, {.scale = 1, .zero_point = 0}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {4}, {.scale = 1e-1, .zero_point = 5}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI16, ElementType::kBF16>(
+      {4}, {.scale = 1e-2, .zero_point = -5}, {-2.22, -1.11, 0, 1.11, 2.22});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {4}, {.scale = 1, .zero_point = 0}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {4}, {.scale = 1e-1, .zero_point = -10}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI16, ElementType::kF16>(
+      {4}, {.scale = 1e-2, .zero_point = 10}, {-2.22, -1.11, 0, 1.11, 2.22});
+
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {4}, {.scale = 1, .zero_point = +7}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {4}, {.scale = 1e-1, .zero_point = -7}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {4}, {.scale = 1e-2, .zero_point = 0}, {-2.22, -1.11, 0, 1.11, 2.22});
+  test<ElementType::kSI32, ElementType::kBF16>(
+      {4}, {.scale = 1e-3, .zero_point = 0}, {-2.222, -1.111, 0, 1.111, 2.222});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {4}, {.scale = 1, .zero_point = +7}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {4}, {.scale = 1e-1, .zero_point = -7}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {4}, {.scale = 1e-2, .zero_point = 10}, {-2.22, -1.11, 0, 1.11, 2.22});
+  test<ElementType::kSI32, ElementType::kF16>(
+      {4}, {.scale = 1e-3, .zero_point = -0},
+      {-2.222, -1.111, 0, 1.111, 2.222});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {4}, {.scale = 1, .zero_point = +7}, {-2, -1, 0, 1, 2});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {4}, {.scale = 1e-1, .zero_point = -7}, {-2.2, -1.1, 0, 1.1, 2.2});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {4}, {.scale = 1e-2, .zero_point = 10}, {-2.22, -1.11, 0, 1.11, 2.22});
+  test<ElementType::kSI32, ElementType::kF32>(
+      {4}, {.scale = 1e-3, .zero_point = -0},
+      {-2.222, -1.111, 0, 1.111, 2.222});
+}
+
+}  // namespace testing
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/test/util.h

@@ -0,0 +1,54 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_UTIL_H_
+
+#include <optional>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace testing {
+
+template <ElementType storage_type, ElementType expressed_type>
+std::vector<typename Storage<storage_type>::Type> QuantizeVector(
+    const std::vector<typename Storage<expressed_type>::Type>& input,
+    const QuantizedParameter& quantized_parameter) {
+  std::vector<typename Storage<storage_type>::Type> result;
+  typename Storage<expressed_type>::Type scale_inv =
+      1.0 / quantized_parameter.scale;
+  for (auto x : input) {
+    auto q = QuantizePartial<storage_type, expressed_type>(
+        x, scale_inv, quantized_parameter.zero_point);
+    result.push_back(q);
+  }
+  CHECK_OK(CompleteQuantization<storage_type>(  // Crash OK
+      result.data(), result.size(),
+      /* storage_min */ std::nullopt,
+      /* storage_min */ std::nullopt));
+  return result;
+}
+
+}  // namespace testing
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_TEST_UTIL_H_