chore: import upstream snapshot with attribution

tensorflow/lite/experimental/shlo/legacy/BUILD

@@ -0,0 +1,70 @@
+# StableHLO Device Profile reference implementation
+
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [":__subpackages__"],
+)
+
+cc_library(
+    name = "shlo",
+    srcs = [
+        "src/broadcast_in_dim.cc",
+        "src/clamp.cc",
+        "src/compare.cc",
+        "src/concatenate.cc",
+        "src/dispatch.h",
+        "src/elementwise_binary.cc",
+        "src/elementwise_unary.cc",
+        "src/iota.cc",
+        "src/is_finite.cc",
+        "src/select.cc",
+        "src/shlo.cc",
+        "src/uniform_dequantize_quantize.cc",
+    ],
+    hdrs = [
+        "include/shlo.h",
+        "src/storage.h",
+        "src/util.h",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":float",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "debug",
+    srcs = [
+        "src/debug.cc",
+    ],
+    hdrs = [
+        "src/debug.h",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":float",
+        ":shlo",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "float",
+    srcs = [
+    ],
+    hdrs = [
+        "src/bf16.h",
+        "src/f16.h",
+        "src/has_keyword.h",
+    ],
+    tags = ["no_oss"],
+    deps = [
+    ],
+)

tensorflow/lite/experimental/shlo/legacy/bench/BUILD

@@ -0,0 +1,51 @@
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
+
+cc_library(
+    name = "util",
+    srcs = [
+    ],
+    hdrs = [
+        "util.h",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/lite/experimental/shlo/legacy:float",
+    ],
+)
+
+cc_binary(
+    name = "shlo_benchmark",
+    srcs = [
+        "shlo_benchmark.cc",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:shlo",
+        "//tensorflow/lite/experimental/shlo/legacy/test:util",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_benchmark//:benchmark",
+    ],
+)
+
+cc_binary(
+    name = "xnn_benchmark",
+    srcs = [
+        "xnn_benchmark.cc",
+    ],
+    data = [
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":util",
+        "//tensorflow/lite/experimental/shlo/legacy:float",
+        "@XNNPACK",
+        "@com_google_absl//absl/log",
+        "@com_google_benchmark//:benchmark",
+    ],
+)

tensorflow/lite/experimental/shlo/legacy/bench/shlo_benchmark.cc

@@ -0,0 +1,180 @@
+/* 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 <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "benchmark/benchmark.h"  // from @com_google_benchmark
+#include "tensorflow/lite/experimental/shlo/legacy/bench/util.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/test/util.h"
+
+namespace stablehlo {
+namespace benchmark {
+namespace {
+
+template <absl::Status (*op)(const Tensor&, Tensor&), ElementType element_type,
+          int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+
+  using ST = typename Storage<element_type>::Type;
+  auto operand_values = GenerateRandomVector<ST>(size);
+  decltype(operand_values) result_values(operand_values.size());
+
+  Tensor operand(TensorType(Shape(shape), element_type),
+                 std::data(operand_values));
+  Tensor result(TensorType(Shape(shape), element_type),
+                std::data(result_values));
+
+  for (auto _ : state) {
+    auto res = op(operand, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const Tensor&, const Tensor&, Tensor&),
+          ElementType element_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+
+  using ST = typename Storage<element_type>::Type;
+  auto lhs_values = GenerateRandomVector<ST>(size);
+  auto rhs_values = GenerateRandomVector<ST>(size);
+  decltype(rhs_values) result_values(rhs_values.size());
+
+  Tensor lhs(TensorType(Shape(shape), element_type), std::data(lhs_values));
+  Tensor rhs(TensorType(Shape(shape), element_type), std::data(rhs_values));
+  Tensor result(TensorType(Shape(shape), element_type),
+                std::data(result_values));
+
+  for (auto _ : state) {
+    auto res = op(lhs, rhs, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const QuantizedTensor&, QuantizedTensor&),
+          ElementType storage_type, ElementType expressed_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+  QuantizedParameter quantized_parameter = {.scale = 0.1, .zero_point = 0};
+
+  using ET = typename Storage<expressed_type>::Type;
+
+  auto operand_values = GenerateRandomVector<ET>(size);
+
+  auto operand_quant_values =
+      testing::QuantizeVector<storage_type, expressed_type>(
+          operand_values, quantized_parameter);
+  decltype(operand_quant_values) result_quant_values(
+      operand_quant_values.size());
+
+  QuantizedTensorElementType element_type(storage_type, expressed_type,
+                                          std::move(quantized_parameter));
+
+  QuantizedTensor operand(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      operand_quant_values.data());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  for (auto _ : state) {
+    auto res = op(operand, result);
+    CHECK_OK(res);
+  }
+}
+
+template <absl::Status (*op)(const QuantizedTensor&, const QuantizedTensor&,
+                             QuantizedTensor&),
+          ElementType storage_type, ElementType expressed_type, int size>
+void BM_SHLO(::benchmark::State& state) {
+  Shape shape = {size};
+  QuantizedParameter quantized_parameter = {.scale = 0.1, .zero_point = 0};
+
+  using ET = typename Storage<expressed_type>::Type;
+
+  auto lhs_values = GenerateRandomVector<ET>(size);
+  auto rhs_values = GenerateRandomVector<ET>(size);
+
+  auto lhs_quant_values = testing::QuantizeVector<storage_type, expressed_type>(
+      lhs_values, quantized_parameter);
+  auto rhs_quant_values = testing::QuantizeVector<storage_type, expressed_type>(
+      rhs_values, quantized_parameter);
+  decltype(rhs_quant_values) result_quant_values(rhs_quant_values.size());
+
+  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());
+  QuantizedTensor result(
+      QuantizedTensorType(Shape(shape),
+                          QuantizedTensorElementType(element_type)),
+      result_quant_values.data());
+
+  for (auto _ : state) {
+    auto res = op(lhs, rhs, result);
+    CHECK_OK(res);
+  }
+}
+
+#define BENCHMARK_OP_HELPER(Op, ElementType)    \
+  BENCHMARK(BM_SHLO<Op, ElementType, 8 * KB>);  \
+  BENCHMARK(BM_SHLO<Op, ElementType, 16 * KB>); \
+  BENCHMARK(BM_SHLO<Op, ElementType, 32 * KB>); \
+  BENCHMARK(BM_SHLO<Op, ElementType, 64 * KB>);
+
+#define BENCHMARK_QOP_HELPER(Op, StorageType, ExpressedType)   \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 8 * KB>);  \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 16 * KB>); \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 32 * KB>); \
+  BENCHMARK(BM_SHLO<Op, StorageType, ExpressedType, 64 * KB>);
+
+#define BENCHMARK_OP(Op)                                            \
+  BENCHMARK_OP_HELPER(Op, ElementType::kF32);                       \
+  BENCHMARK_OP_HELPER(Op, ElementType::kF16);                       \
+  BENCHMARK_OP_HELPER(Op, ElementType::kBF16);                      \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kF32);   \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kF32);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kF32);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kF16);   \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI8, ElementType::kBF16);  \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI16, ElementType::kBF16); \
+  BENCHMARK_QOP_HELPER(Op, ElementType::kSI32, ElementType::kBF16);
+
+BENCHMARK_OP(Abs);
+BENCHMARK_OP(Add);
+
+}  // namespace
+}  // namespace benchmark
+}  // namespace stablehlo
+
+// Run the benchmark
+BENCHMARK_MAIN();

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

@@ -0,0 +1,64 @@
+/* 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_BENCH_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_BENCH_UTIL_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <limits>
+#include <random>
+#include <type_traits>
+#include <vector>
+
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+
+namespace stablehlo {
+
+namespace benchmark {
+
+static constexpr auto KB = 1024;
+
+template <typename Number>
+std::vector<Number> GenerateRandomVector(
+    size_t size, Number min = std::numeric_limits<Number>::min(),
+    Number max = std::numeric_limits<Number>::max()) {
+  std::vector<Number> data(size);
+  if constexpr (std::is_integral_v<Number>) {
+    static std::uniform_int_distribution<Number> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else if constexpr (std::is_floating_point_v<Number>) {
+    static std::uniform_real_distribution<Number> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else {
+    static_assert(std::is_same_v<Number, BF16> or std::is_same_v<Number, F16>);
+    static std::uniform_real_distribution<float> distribution(min, max);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  }
+  return data;
+}
+
+}  // namespace benchmark
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_BENCH_UTIL_H_

tensorflow/lite/experimental/shlo/legacy/bench/xnn_benchmark.cc

@@ -0,0 +1,241 @@
+/* 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 <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#include "xnnpack.h"  // from @XNNPACK
+#include "absl/log/log.h"
+#include "benchmark/benchmark.h"  // from @com_google_benchmark
+#include "tensorflow/lite/experimental/shlo/legacy/bench/util.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+
+namespace stablehlo {
+namespace benchmark {
+namespace {
+
+template <xnn_datatype datatype>
+struct Storage;
+
+template <>
+struct Storage<xnn_datatype_fp32> {
+  using Type = float;
+};
+
+template <>
+struct Storage<xnn_datatype_fp16> {
+  using Type = F16;
+};
+
+size_t size_in_bytes(xnn_datatype datatype) {
+  switch (datatype) {
+    case xnn_datatype_qint32:
+    case xnn_datatype_qcint32:
+    case xnn_datatype_fp32:
+      return 4;
+    case xnn_datatype_fp16:
+      return 2;
+    case xnn_datatype_qint8:
+    case xnn_datatype_quint8:
+    case xnn_datatype_qcint8:
+      return 1;
+    default:
+      // Crash OK
+      LOG(FATAL) << "Unexpected datatype: " << static_cast<int>(datatype);
+  }
+}
+
+size_t tensor_size_in_bytes(xnn_datatype datatype,
+                            const std::vector<size_t>& shape) {
+  size_t size = size_in_bytes(datatype);
+  for (auto s : shape) {
+    size *= s;
+  }
+  return size;
+}
+
+using UnaryOp = xnn_status (*)(xnn_subgraph_t subgraph, uint32_t input_id,
+                               uint32_t output_id, uint32_t flags);
+
+using BinaryOp = xnn_status (*)(xnn_subgraph_t subgraph, uint32_t input1_id,
+                                uint32_t input2_id, uint32_t output_id,
+                                uint32_t flags);
+
+using ClampedBinaryOp = xnn_status (*)(xnn_subgraph_t subgraph,
+                                       float output_min, float output_max,
+                                       uint32_t input1_id, uint32_t input2_id,
+                                       uint32_t output_id, uint32_t flags);
+
+template <xnn_datatype datatype, typename Op>
+void BM_XNN_HELPER(::benchmark::State& state, Op op, size_t size) {
+  size_t num_operands;
+  if constexpr (std::is_same_v<Op, UnaryOp>) {
+    num_operands = 1;
+  } else {
+    static_assert(std::is_same_v<Op, BinaryOp> or
+                  std::is_same_v<Op, ClampedBinaryOp>);
+    num_operands = 2;
+  }
+
+  std::vector<size_t> shape = {size};
+
+  using ST = typename Storage<datatype>::Type;
+
+  std::vector<std::vector<ST>> operand_values;
+  operand_values.reserve(num_operands);
+  for (auto i = 0; i < num_operands; ++i) {
+    operand_values.emplace_back(GenerateRandomVector<ST>(size));
+  }
+
+  auto size_in_bytes = tensor_size_in_bytes(datatype, shape);
+  std::vector<uint8_t> result_values(size_in_bytes);
+
+  std::vector<xnn_external_value> external_values;
+  external_values.reserve(num_operands + 1);
+  for (auto i = 0; i < num_operands; ++i) {
+    external_values.push_back(
+        {.id = static_cast<uint32_t>(i), .data = operand_values[i].data()});
+  }
+  external_values.push_back({.id = static_cast<uint32_t>(num_operands),
+                             .data = result_values.data()});
+
+  xnn_status status = xnn_initialize(/*allocator=*/nullptr);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_workspace_t workspace;
+  status = xnn_create_workspace(&workspace);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_subgraph_t subgraph = nullptr;
+  auto max_external_value_id =
+      std::max_element(external_values.begin(), external_values.end(),
+                       [](auto x, auto y) { return x.id < y.id; })
+          ->id;
+  status =
+      xnn_create_subgraph(1 + max_external_value_id, /*flags=*/0, &subgraph);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  for (auto i = 0; i < num_operands; ++i) {
+    uint32_t operand_id;
+    status = xnn_define_tensor_value(
+        subgraph, datatype, shape.size(), shape.data(),
+        /* data */ nullptr, /* xnn_external_id */ i,
+        /* flags */ XNN_VALUE_FLAG_EXTERNAL_INPUT, &operand_id);
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+  }
+
+  uint32_t result_id;
+  status = xnn_define_tensor_value(
+      subgraph, datatype, shape.size(), shape.data(),
+      /* data */ nullptr, /* xnn_external_id */ num_operands,
+      /* flags */ XNN_VALUE_FLAG_EXTERNAL_OUTPUT, &result_id);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  if constexpr (std::is_same_v<Op, UnaryOp>) {
+    status = op(subgraph, 0, 1, /* flags */ 0);
+  } else if constexpr (std::is_same_v<Op, BinaryOp>) {
+    status = op(subgraph, 0, 1, 2, /* flags */ 0);
+  } else {
+    static_assert(std::is_same_v<Op, ClampedBinaryOp>);
+    status = op(subgraph, /* output_min */ std::numeric_limits<float>::min(),
+                /* output_max */ std::numeric_limits<float>::max(), 0, 1, 2,
+                /* flags */ 0);
+  }
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  xnn_runtime_t runtime = nullptr;
+  status = xnn_create_runtime_v4(subgraph,
+                                 /*weights_buffer_cache*/ nullptr, workspace,
+                                 /*threadpool*/ nullptr,
+                                 /*flags*/ 0, &runtime);
+  if (status != xnn_status_success) {
+    // Crash OK
+    LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+  }
+
+  for (auto _ : state) {
+    status = xnn_setup_runtime(runtime, external_values.size(),
+                               external_values.data());
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+    status = xnn_invoke_runtime(runtime);
+    if (status != xnn_status_success) {
+      // Crash OK
+      LOG(FATAL) << "Failed to invoke XNNPACK runtime: " << status;
+    }
+  }
+
+  xnn_delete_runtime(runtime);
+  xnn_delete_subgraph(subgraph);
+  xnn_release_workspace(workspace);
+  xnn_deinitialize();
+}
+
+template <UnaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+template <BinaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+template <ClampedBinaryOp op, xnn_datatype datatype, size_t size>
+void BM_XNN(::benchmark::State& state) {
+  BM_XNN_HELPER<datatype>(state, op, size);
+}
+
+#define BENCHMARK_OP_HELPER(Op, DataType)   \
+  BENCHMARK(BM_XNN<Op, DataType, 8 * KB>);  \
+  BENCHMARK(BM_XNN<Op, DataType, 16 * KB>); \
+  BENCHMARK(BM_XNN<Op, DataType, 32 * KB>); \
+  BENCHMARK(BM_XNN<Op, DataType, 64 * KB>);
+
+#define BENCHMARK_OP(Op) BENCHMARK_OP_HELPER(Op, xnn_datatype_fp32);
+//  BENCHMARK_OP_HELPER(Op, xnn_datatype_fp16);
+
+BENCHMARK_OP(xnn_define_abs);
+BENCHMARK_OP(xnn_define_add2);
+
+}  // namespace
+}  // namespace benchmark
+}  // namespace stablehlo
+
+// Run the benchmark
+BENCHMARK_MAIN();

tensorflow/lite/experimental/shlo/legacy/include/shlo.h

@@ -0,0 +1,422 @@
+/* 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_INCLUDE_SHLO_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_INCLUDE_SHLO_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+
+namespace stablehlo {
+
+enum class ElementType { kUnknown, kI1, kSI8, kSI16, kSI32, kBF16, kF16, kF32 };
+
+using DimensionSize = size_t;
+using Axes = std::vector<size_t>;
+using Dims = absl::Span<const DimensionSize>;
+
+/*
+  A tensor shape represents non-negative dimension sizes in the ascending order
+  of the corresponding dimensions (which are also called axes) numbered from 0
+  to R-1. The number of dimensions R is called rank. For example,
+  tensor<2x3xf32> is a tensor type with shape 2x3 and element type f32. It has
+  two dimensions (or, in other words, two axes) - 0th dimension and 1st
+  dimension - whose sizes are 2 and 3. Its rank is 2.
+*/
+class Shape {
+ public:
+  // Scalar shape, with rank 0
+  Shape() = default;
+  // Tensor shape, with rank > 0
+  explicit Shape(std::vector<DimensionSize>&& dims) : dims_(std::move(dims)) {}
+  Shape(std::initializer_list<DimensionSize>&& dims) : dims_(std::move(dims)) {}
+  bool operator==(const Shape& other) const { return dims_ == other.dims_; }
+  bool operator!=(const Shape& other) const { return !(*this == other); }
+  size_t rank() const { return dims_.size(); }
+  DimensionSize dim(size_t idx) const { return dims_[idx]; }
+  Axes axes() const;
+  Dims dims() const { return dims_; }
+  size_t num_elements() const;
+
+ private:
+  std::vector<DimensionSize> dims_;
+};
+
+class TensorType {
+ public:
+  TensorType(Shape&& shape, ElementType tensor_element_type)
+      : shape_(std::move(shape)), tensor_element_type_(tensor_element_type) {}
+  bool operator==(const TensorType& other) const {
+    return shape_ == other.shape_ and
+           tensor_element_type_ == other.tensor_element_type_;
+  }
+  bool operator!=(const TensorType& other) const { return !(*this == other); }
+  const Shape& shape() const { return shape_; }
+  ElementType element_type() const { return tensor_element_type_; }
+  size_t num_bytes() const;
+
+ private:
+  Shape shape_;
+  ElementType tensor_element_type_;
+};
+
+// Tensor layout has is not part of the StableHLO standard. At this time we
+// support only strides for the minor dimension.
+class Layout {
+ public:
+  explicit Layout(size_t minor_dim_stride = 1)
+      : minor_dim_stride_(minor_dim_stride) {}
+  bool has_strides() const { return minor_dim_stride_ > 1; }
+  size_t minor_dim_stride() const { return minor_dim_stride_; }
+
+ private:
+  size_t minor_dim_stride_ = 1;
+};
+
+class Tensor {
+ public:
+  Tensor(TensorType&& type, void* buffer, Layout&& layout = Layout())
+      : type_(std::move(type)), buffer_(buffer), layout_(std::move(layout)) {}
+  bool operator==(const Tensor& other) const;
+  bool operator!=(const Tensor& other) const { return !(*this == other); }
+  const TensorType& type() const { return type_; }
+  void* buffer() { return buffer_; }
+  const void* buffer() const { return buffer_; }
+  const Layout& layout() const { return layout_; }
+  const Shape& shape() const { return type_.shape(); }
+  ElementType element_type() const { return type_.element_type(); }
+  TensorType baseline_type() const { return type_; }
+  ElementType baseline_element_type() const {
+    return baseline_type().element_type();
+  }
+  size_t rank() const { return shape().rank(); }
+  DimensionSize dim(size_t idx) const { return shape().dim(idx); }
+  Axes axes() const { return shape().axes(); }
+  Dims dims() const { return shape().dims(); }
+  size_t num_elements() const { return shape().num_elements(); }
+  size_t num_bytes() const { return type_.num_bytes(); }
+  bool is_per_tensor_quantized() const { return false; }
+  bool is_per_axis_quantized() const { return false; }
+
+ private:
+  TensorType type_;
+  void* buffer_;
+  Layout layout_;
+};
+
+struct QuantizedParameter {
+  float scale;
+  int32_t zero_point;
+  bool operator==(const QuantizedParameter& other) const {
+    return scale == other.scale and zero_point == other.zero_point;
+  }
+  bool operator!=(const QuantizedParameter& other) const {
+    return !(*this == other);
+  }
+};
+
+class QuantizedTensorElementType {
+ public:
+  // Constructor for per-tensor quantization.
+  QuantizedTensorElementType(ElementType storage_type,
+                             ElementType expressed_type,
+                             QuantizedParameter&& quantized_parameter,
+                             std::optional<int32_t> storage_min = std::nullopt,
+                             std::optional<int32_t> storage_max = std::nullopt)
+      : storage_type_(storage_type), expressed_type_(expressed_type) {
+    parameters_.emplace_back(std::move(quantized_parameter));
+  }
+
+  // Constructor for per-axis quantization.
+  QuantizedTensorElementType(ElementType storage_type,
+                             ElementType expressed_type,
+                             std::vector<QuantizedParameter>&& parameters,
+                             std::optional<int32_t> storage_min = std::nullopt,
+                             std::optional<int32_t> storage_max = std::nullopt)
+      : storage_type_(storage_type),
+        expressed_type_(expressed_type),
+        parameters_(std::move(parameters)) {}
+
+  ElementType storage_type() const { return storage_type_; }
+  ElementType expressed_type() const { return expressed_type_; }
+  std::optional<int32_t> storage_min() const { return storage_min_; }
+  std::optional<int32_t> storage_max() const { return storage_max_; }
+  std::optional<DimensionSize> quantized_dimension() const {
+    return quantized_dimension_;
+  }
+  const QuantizedParameter& parameters(size_t index) const {
+    return parameters_[index];
+  }
+  size_t num_parameters() const { return parameters_.size(); }
+  bool operator==(const QuantizedTensorElementType& other) const {
+    return storage_type_ == other.storage_type_ and
+           expressed_type_ == other.expressed_type_ and
+           storage_min_ == other.storage_min_ and
+           storage_max_ == other.storage_max_ and
+           quantized_dimension_ == other.quantized_dimension_ and
+           parameters_ == other.parameters_;
+  }
+  bool operator!=(const QuantizedTensorElementType& other) const {
+    return !(*this == other);
+  }
+  bool is_per_tensor_quantized() const { return !quantized_dimension_; }
+  bool is_per_axis_quantized() const { return !is_per_tensor_quantized(); }
+
+ private:
+  ElementType storage_type_;
+  ElementType expressed_type_;
+  std::optional<int32_t> storage_min_;
+  std::optional<int32_t> storage_max_;
+  std::optional<DimensionSize> quantized_dimension_;
+  std::vector<QuantizedParameter> parameters_;
+};
+
+class QuantizedTensorType {
+ public:
+  QuantizedTensorType(
+      Shape&& shape, QuantizedTensorElementType&& quantized_tensor_element_type)
+      : shape_(std::move(shape)),
+        quantized_tensor_element_type_(
+            std::move(quantized_tensor_element_type)) {}
+  const Shape& shape() const { return shape_; }
+  const QuantizedTensorElementType& element_type() const {
+    return quantized_tensor_element_type_;
+  }
+  bool operator==(const QuantizedTensorType& other) const {
+    return shape_ == other.shape_ and quantized_tensor_element_type_ ==
+                                          other.quantized_tensor_element_type_;
+  }
+  bool operator!=(const QuantizedTensorType& other) const {
+    return !(*this == other);
+  }
+  size_t num_bytes() const;
+
+ private:
+  Shape shape_;
+  QuantizedTensorElementType quantized_tensor_element_type_;
+};
+
+class QuantizedTensor {
+ public:
+  QuantizedTensor(QuantizedTensorType&& type, void* buffer,
+                  Layout&& layout = Layout())
+      : type_(std::move(type)), buffer_(buffer), layout_(std::move(layout)) {}
+  bool operator==(const QuantizedTensor& other) const;
+  bool operator!=(const QuantizedTensor& other) const {
+    return !(*this == other);
+  }
+  const QuantizedTensorType& type() const { return type_; }
+  void* buffer() { return buffer_; }
+  const void* buffer() const { return buffer_; }
+  const Layout& layout() const { return layout_; }
+  const Shape& shape() const { return type_.shape(); }
+  ElementType storage_type() const { return element_type().storage_type(); }
+  ElementType expressed_type() const { return element_type().expressed_type(); }
+  const QuantizedTensorElementType& element_type() const {
+    return type_.element_type();
+  }
+  QuantizedTensorType baseline_type() const;
+  QuantizedTensorElementType baseline_element_type() const {
+    return baseline_type().element_type();
+  }
+  size_t rank() const { return shape().rank(); }
+  DimensionSize dim(size_t idx) const { return shape().dim(idx); }
+  Axes axes() const { return shape().axes(); }
+  Dims dims() const { return shape().dims(); }
+  size_t num_elements() const { return shape().num_elements(); }
+  size_t num_bytes() const { return type_.num_bytes(); }
+  bool is_per_tensor_quantized() const {
+    return element_type().is_per_tensor_quantized();
+  }
+  bool is_per_axis_quantized() const {
+    return element_type().is_per_axis_quantized();
+  }
+  auto quantized_dimension() const {
+    return element_type().quantized_dimension();
+  }
+  auto storage_min() const { return element_type().storage_min(); }
+  auto storage_max() const { return element_type().storage_max(); }
+  auto scales(size_t idx) const { return element_type().parameters(idx).scale; }
+  auto zero_points(size_t idx) const {
+    return element_type().parameters(idx).zero_point;
+  }
+
+ private:
+  QuantizedTensorType type_;
+  void* buffer_;
+  Layout layout_;
+};
+
+inline bool IsSignedInteger(ElementType element_type) {
+  return element_type == ElementType::kSI8 ||
+         element_type == ElementType::kSI16 ||
+         element_type == ElementType::kSI32;
+}
+
+inline bool IsUnsignedInteger(ElementType element_type) { return false; }
+
+inline bool IsBoolean(ElementType element_type) {
+  return element_type == ElementType::kI1;
+}
+
+inline bool IsFloat(ElementType element_type) {
+  return element_type == ElementType::kBF16 ||
+         element_type == ElementType::kF16 || element_type == ElementType::kF32;
+}
+
+inline bool IsSignedInteger(const QuantizedTensorElementType& element_type) {
+  return IsSignedInteger(element_type.expressed_type());
+}
+inline bool IsUnsignedInteger(const QuantizedTensorElementType& element_type) {
+  return IsUnsignedInteger(element_type.expressed_type());
+}
+inline bool IsBoolean(const QuantizedTensorElementType& element_type) {
+  return IsBoolean(element_type.expressed_type());
+}
+inline bool IsFloat(const QuantizedTensorElementType& element_type) {
+  return IsFloat(element_type.expressed_type());
+}
+
+enum class ComparisonDirection { kEQ, kNE, kGE, kGT, kLE, kLT };
+enum class CompareType {
+  kFloat,
+  kTotalOrder /* Unsupported */,
+  kSigned,
+  kUnsigned
+};
+
+// /////////////////////////////////////////////////////////////////////////////
+
+absl::Status Abs(const Tensor& operand, Tensor& result);
+absl::Status Abs(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Add(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Add(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                 QuantizedTensor& result);
+absl::Status And(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Atan2(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Atan2(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                   QuantizedTensor& result);
+absl::Status BroadcastInDim(
+    const Tensor& operand, absl::Span<const DimensionSize> broadcast_dimensions,
+    Tensor& result);
+absl::Status BroadcastInDim(
+    const QuantizedTensor& operand,
+    absl::Span<const DimensionSize> broadcast_dimensions,
+    QuantizedTensor& result);
+absl::Status Clamp(const Tensor& min, const Tensor& operand, const Tensor& max,
+                   Tensor& result);
+absl::Status Clamp(const QuantizedTensor& min, const QuantizedTensor& operand,
+                   const QuantizedTensor& max, QuantizedTensor& result);
+absl::Status Cbrt(const Tensor& operand, Tensor& result);
+absl::Status Cbrt(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Ceil(const Tensor& operand, Tensor& result);
+absl::Status Ceil(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Compare(const Tensor& lhs, const Tensor& rhs,
+                     ComparisonDirection comparison_direction,
+                     CompareType compare_type, Tensor& result);
+absl::Status Compare(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     ComparisonDirection comparison_direction,
+                     CompareType compare_type, Tensor& result);
+absl::Status Concatenate(absl::Span<const Tensor*> inputs,
+                         DimensionSize dimension, Tensor& result);
+absl::Status Concatenate(absl::Span<const QuantizedTensor*> inputs,
+                         DimensionSize dimension, QuantizedTensor& result);
+absl::Status Cosine(const Tensor& operand, Tensor& result);
+absl::Status Cosine(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status CountLeadingZeros(const Tensor& operand, Tensor& result);
+absl::Status Divide(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Divide(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                    QuantizedTensor& result);
+absl::Status Exponential(const Tensor& operand, Tensor& result);
+absl::Status Exponential(const QuantizedTensor& operand,
+                         QuantizedTensor& result);
+absl::Status ExponentialMinusOne(const Tensor& operand, Tensor& result);
+absl::Status ExponentialMinusOne(const QuantizedTensor& operand,
+                                 QuantizedTensor& result);
+absl::Status Floor(const Tensor& operand, Tensor& result);
+absl::Status Floor(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Iota(DimensionSize iota_dimension, Tensor& result);
+absl::Status Iota(DimensionSize iota_dimension, QuantizedTensor& result);
+absl::Status IsFinite(const Tensor& operand, Tensor& result);
+absl::Status IsFinite(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Log(const Tensor& operand, Tensor& result);
+absl::Status Log(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status LogPlusOne(const Tensor& operand, Tensor& result);
+absl::Status LogPlusOne(const QuantizedTensor& operand,
+                        QuantizedTensor& result);
+absl::Status Logistic(const Tensor& operand, Tensor& result);
+absl::Status Logistic(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Maximum(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Maximum(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     QuantizedTensor& result);
+absl::Status Minimum(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Minimum(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     QuantizedTensor& result);
+absl::Status Multiply(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Multiply(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                      QuantizedTensor& result);
+absl::Status Negate(const Tensor& operand, Tensor& result);
+absl::Status Negate(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Not(const Tensor& operand, Tensor& result);
+absl::Status Or(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Popcnt(const Tensor& operand, Tensor& result);
+absl::Status Power(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Power(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                   QuantizedTensor& result);
+absl::Status Remainder(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Remainder(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                       QuantizedTensor& result);
+absl::Status RoundNearestAfz(const Tensor& operand, Tensor& result);
+absl::Status RoundNearestAfz(const QuantizedTensor& operand,
+                             QuantizedTensor& result);
+absl::Status RoundNearestEven(const Tensor& operand, Tensor& result);
+absl::Status RoundNearestEven(const QuantizedTensor& operand,
+                              QuantizedTensor& result);
+absl::Status Rsqrt(const Tensor& operand, Tensor& result);
+absl::Status Rsqrt(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Select(const Tensor& pred, const Tensor& on_true,
+                    const Tensor& on_false, Tensor& result);
+absl::Status Select(const Tensor& pred, const QuantizedTensor& on_true,
+                    const QuantizedTensor& on_false, QuantizedTensor& result);
+absl::Status ShiftLeft(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status ShiftRightArithmetic(const Tensor& lhs, const Tensor& rhs,
+                                  Tensor& result);
+absl::Status ShiftRightLogical(const Tensor& lhs, const Tensor& rhs,
+                               Tensor& result);
+absl::Status Sign(const Tensor& operand, Tensor& result);
+absl::Status Sign(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Sine(const Tensor& operand, Tensor& result);
+absl::Status Sine(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Subtract(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+absl::Status Subtract(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                      QuantizedTensor& result);
+absl::Status Sqrt(const Tensor& operand, Tensor& result);
+absl::Status Sqrt(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status Tanh(const Tensor& operand, Tensor& result);
+absl::Status Tanh(const QuantizedTensor& operand, QuantizedTensor& result);
+absl::Status UniformDequantize(const QuantizedTensor& operand, Tensor& result);
+absl::Status UniformQuantize(const Tensor& operand, QuantizedTensor& result);
+absl::Status Xor(const Tensor& lhs, const Tensor& rhs, Tensor& result);
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_INCLUDE_SHLO_H_

tensorflow/lite/experimental/shlo/legacy/src/bf16.h

@@ -0,0 +1,129 @@
+/* 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_SRC_BF16_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_BF16_H_
+
+#include "tensorflow/lite/experimental/shlo/legacy/src/has_keyword.h"
+
+#if defined(__STDCPP_BFLOAT16_T__)
+#include <stdfloat>
+namespace stablehlo {
+using BF16 = bfloat16_t;
+}  // namespace stablehlo
+
+#elif __has_keyword(__bf16) && __x86_64__
+namespace stablehlo {
+// On x86 the compiler is able to generate code for __bf16 operations.
+using BF16 = __bf16;
+}  // namespace stablehlo
+
+#elif __has_keyword(__bf16) && __aarch64__
+#include <cmath>
+#include <cstdint>
+
+namespace stablehlo {
+
+// On arm64 the compiler is not yet able to generate code for __bf16
+// operations. Therefore, we resort to a software-based implementation of BF16
+// based on promoting ops to float.
+class BF16 {
+ public:
+  BF16(float f = 0.0f) {
+    if (std::isnan(f)) {
+      // If the value is a NaN, squash it to a NaN with the msb of the
+      // mantissa. This avoids that after the truncation below we don't end up
+      // with an inf.
+      value_ = std::signbit(f) ? 0xFFC0 : 0x7FC0;
+    } else {
+      // Fast rounding algorithm that rounds a half value to nearest even. This
+      // reduces expected error when we convert a large number of floats.
+      uint32_t input = *reinterpret_cast<const uint32_t*>(&f);
+
+      // Least significant bit of resulting bfloat.
+      uint32_t lsb = (input >> 16) & 1;
+      uint32_t rounding_bias = 0x7fff + lsb;
+      input += rounding_bias;
+
+      value_ = static_cast<uint16_t>(input >> 16u);
+    }
+  }
+
+  BF16& operator=(BF16 other) {
+    value_ = other.value_;
+    return *this;
+  }
+
+  bool operator==(BF16 other) const { return value_ == other.value_; }
+  bool operator!=(BF16 other) const { return !(*this == other); }
+
+  operator float() const {
+    uint32_t tmp = value_ << 16;
+    return *reinterpret_cast<float*>(&tmp);
+  }
+
+  BF16 operator-() const { return BF16(-static_cast<float>(*this)); }
+
+  BF16& operator+=(BF16 other) {
+    value_ = BF16(static_cast<float>(*this) + static_cast<float>(other)).value_;
+    return *this;
+  }
+
+  BF16& operator-=(BF16 other) {
+    value_ = BF16(static_cast<float>(*this) - static_cast<float>(other)).value_;
+    return *this;
+  }
+
+  BF16& operator*=(BF16 other) {
+    value_ = BF16(static_cast<float>(*this) * static_cast<float>(other)).value_;
+    return *this;
+  }
+
+  BF16& operator/=(BF16 other) {
+    value_ = BF16(static_cast<float>(*this) / static_cast<float>(other)).value_;
+    return *this;
+  }
+
+ private:
+  uint16_t value_;
+};
+
+inline BF16 operator+(BF16 x, BF16 y) {
+  x += y;
+  return x;
+}
+
+inline BF16 operator-(BF16 x, BF16 y) {
+  x -= y;
+  return x;
+}
+
+inline BF16 operator*(BF16 x, BF16 y) {
+  x *= y;
+  return x;
+}
+
+inline BF16 operator/(BF16 x, BF16 y) {
+  x /= y;
+  return x;
+}
+
+}  // namespace stablehlo
+
+#else
+#error Type BF16 is not available
+#endif
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_BF16_H_

tensorflow/lite/experimental/shlo/legacy/src/broadcast_in_dim.cc

@@ -0,0 +1,223 @@
+/* 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 <algorithm>
+#include <iterator>
+#include <type_traits>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+bool IsUnique(absl::Span<const DimensionSize> span) {
+  std::vector<DimensionSize> temp(span.begin(), span.end());
+  auto i = std::unique(temp.begin(), temp.end());
+  return std::distance(temp.begin(), i) == span.size();
+}
+
+template <typename Value>
+absl::Status CheckParameters(
+    const Value& operand, absl::Span<const DimensionSize> broadcast_dimensions,
+    Value& result) {
+  if (!operand.is_per_axis_quantized()) {
+    if (!(result.element_type() == operand.element_type())) {
+      return absl::InvalidArgumentError(
+          "Constraint violation: element_type(result) = element_type(operand) "
+          "if !is_per_axis_quantized(operand)");
+    }
+  }
+
+  if (!(broadcast_dimensions.size() == operand.rank())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: size(broadcast_dimensions) = rank(operand)");
+  } else if (!(*std::min_element(broadcast_dimensions.begin(),
+                                 broadcast_dimensions.end()) >= 0 and
+               *std::max_element(broadcast_dimensions.begin(),
+                                 broadcast_dimensions.end()) < result.rank())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: 0 <= broadcast_dimensions < rank(result)");
+  } else if (!(IsUnique(broadcast_dimensions))) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: is_unique(broadcast_dimensions)");
+  } else {
+    for (auto d : operand.axes()) {
+      if (!(operand.dim(d) == 1 or
+            operand.dim(d) == result.dim(broadcast_dimensions[d]))) {
+        return absl::InvalidArgumentError(
+            "Constraint violation: dim(operand, d) = 1 or dim(operand, d) = "
+            "dim(result, broadcast_dimensions[d])");
+      }
+    }
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (operand.is_per_axis_quantized()) {
+      if (!(operand.is_per_axis_quantized() and
+            result.storage_type() == operand.storage_type() and
+            result.expressed_type() == operand.expressed_type() and
+            result.storage_min() == operand.storage_min() and
+            result.storage_max() == operand.storage_max())) {
+        return absl::InvalidArgumentError(
+            "Constraint violation: element_type(result) = "
+            "element_type(operand) with exceptions if "
+            "is_per_axis_quantized(operand)");
+      }
+    }
+    if (result.is_per_axis_quantized()) {
+      if (!(*result.quantized_dimension() ==
+            broadcast_dimensions[*operand.quantized_dimension()])) {
+        return absl::InvalidArgumentError(
+            "quantization_dimension(result) = "
+            "broadcast_dimensions[quantization_dimension(operand)]");
+      }
+      if (operand.dim(*operand.quantized_dimension()) == 1) {
+        auto n = result.dim(*result.quantized_dimension());
+        for (auto i = 0; i < n; ++i) {
+          if (!(result.scales(i) == operand.scales(0) and
+                result.zero_points(i) == operand.zero_points(0))) {
+            return absl::InvalidArgumentError(
+                "If dim(operand, quantization_dimension(operand)) = 1, then "
+                "scales(result)[i] = scales(operand)[0] and "
+                "zero_points(result)[i] = zero_points(operand)[0] for i in "
+                "range(dim(result, quantization_dimension(result)))");
+          }
+        }
+      }
+    }
+  }
+
+  if (operand.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status BroadcastInDim(
+    const Value& operand, absl::Span<const DimensionSize> broadcast_dimensions,
+    Value& result) {
+  if (auto check = CheckParameters(operand, broadcast_dimensions, result);
+      !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  auto operand_buffer = operand.buffer();
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != operand.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    TensorIndex operand_index(operand.shape());
+    for (TensorIndexIterator result_index_iter{result.shape()};
+         result_index_iter.has_next(); ++result_index_iter) {
+      for (auto d = 0; d < operand.rank(); ++d) {
+        if (operand.dim(d) == 1) {
+          operand_index.set(d, 0);
+        } else {
+          auto b = broadcast_dimensions[d];
+          operand_index.set(d, (*result_index_iter)[b]);
+        }
+      }
+      auto linearized_operand_index = operand_index.linearize();
+      auto linearized_result_index = result_index_iter->linearize();
+      auto value = S::Get(operand_buffer, linearized_operand_index);
+      S::Set(result_buffer, linearized_result_index, value);
+    }
+
+  } else if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    if (!(operand.is_per_tensor_quantized() and
+          result.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError(
+          "Only per-tensor quantization is currently supported");
+    }
+
+    using ET = typename Storage<expressed_type>::Type;
+
+    const QuantizedParameter& operand_quant_param =
+        operand.type().element_type().parameters(0);
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    TensorIndex operand_index(operand.shape());
+    for (TensorIndexIterator result_index_iter{result.shape()};
+         result_index_iter.has_next(); ++result_index_iter) {
+      for (auto d = 0; d < operand.rank(); ++d) {
+        if (operand.dim(d) == 1) {
+          operand_index.set(d, 0);
+        } else {
+          auto b = broadcast_dimensions[d];
+          operand_index.set(d, (*result_index_iter)[b]);
+        }
+      }
+
+      auto linearized_operand_index = operand_index.linearize();
+      auto linearized_result_index = result_index_iter->linearize();
+
+      auto operand_storage = S::Get(operand_buffer, linearized_operand_index);
+      auto result_storage =
+          DequantizeOpQuantizePartial<storage_type, expressed_type>(
+              operand_storage, operand_quant_param, result_scale_inv,
+              result_quant_param.zero_point, [](auto x) { return x; });
+      S::Set(result_buffer, linearized_result_index, result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status BroadcastInDim(
+    const Tensor& operand, absl::Span<const DimensionSize> broadcast_dimensions,
+    Tensor& result) {
+  DISPATCH_BOOL_INT_FLOAT(BroadcastInDim, result.element_type(), operand,
+                          broadcast_dimensions, result);
+}
+
+absl::Status BroadcastInDim(
+    const QuantizedTensor& operand,
+    absl::Span<const DimensionSize> broadcast_dimensions,
+    QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(BroadcastInDim, result.storage_type(),
+                     result.expressed_type(), operand, broadcast_dimensions,
+                     result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/clamp.cc

@@ -0,0 +1,170 @@
+/* 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 <algorithm>
+#include <cstddef>
+#include <type_traits>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(const Value& min, const Value& operand,
+                             const Value& max, Value& result) {
+  if (!(min.rank() == 0 or min.shape() == operand.shape())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: rank(min) = 0 or shape(min) = shape(operand)");
+  } else if (!(max.rank() == 0 or max.shape() == operand.shape())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: rank(max) = 0 or shape(max) = shape(operand)");
+  } else if (!(min.baseline_element_type() ==
+                   operand.baseline_element_type() and
+               min.baseline_element_type() == max.baseline_element_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_element_type(min) = "
+        "baseline_element_type(operand) = baseline_element_type(max)");
+  } else if (!(operand.baseline_type() == result.baseline_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_type(operand) = baseline_type(result)");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!(min.is_per_tensor_quantized() and max.is_per_tensor_quantized() and
+          operand.is_per_tensor_quantized() and
+          result.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError("Expected per-tensor quantization");
+    }
+  }
+
+  if (operand.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status Clamp(const Value& min, const Value& operand, const Value& max,
+                   Value& result) {
+  if (auto check = CheckParameters(min, operand, max, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  const bool min_is_tensor = (min.rank() > 0);
+  const bool max_is_tensor = (max.rank() > 0);
+  const size_t n = result.num_elements();
+
+  auto operand_buffer = operand.buffer();
+  auto min_buffer = min.buffer();
+  auto max_buffer = max.buffer();
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != result.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    typename S::Type min_value;
+    typename S::Type max_value;
+    for (size_t i = 0; i < n; ++i) {
+      if (min_is_tensor || (i == 0)) {
+        min_value = S::Get(min_buffer, i);
+      }
+      if (max_is_tensor || (i == 0)) {
+        max_value = S::Get(max_buffer, i);
+      }
+      auto operand_value = S::Get(operand_buffer, i);
+      auto result_value =
+          std::min(max_value, std::max(min_value, operand_value));
+      S::Set(result_buffer, i, result_value);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    using ET = typename Storage<expressed_type>::Type;
+
+    const QuantizedParameter& min_quant_param =
+        min.type().element_type().parameters(0);
+    const QuantizedParameter& max_quant_param =
+        max.type().element_type().parameters(0);
+    const QuantizedParameter& operand_quant_param =
+        operand.type().element_type().parameters(0);
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    ET min_expressed;
+    ET max_expressed;
+    for (size_t i = 0; i < n; ++i) {
+      if (min_is_tensor || (i == 0)) {
+        auto min_storage = S::Get(min_buffer, i);
+        min_expressed = Dequantize<storage_type, expressed_type>(
+            min_storage, min_quant_param);
+      }
+      if (max_is_tensor || (i == 0)) {
+        auto max_storage = S::Get(max_buffer, i);
+        max_expressed = Dequantize<storage_type, expressed_type>(
+            max_storage, max_quant_param);
+      }
+
+      auto operand_storage = S::Get(operand_buffer, i);
+      auto result_storage =
+          DequantizeOpQuantizePartial<storage_type, expressed_type>(
+              operand_storage, operand_quant_param, result_scale_inv,
+              result_quant_param.zero_point, [=](auto x) {
+                return std::min(max_expressed, std::max(min_expressed, x));
+              });
+      S::Set(result_buffer, i, result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status Clamp(const Tensor& min, const Tensor& operand, const Tensor& max,
+                   Tensor& result) {
+  DISPATCH_INT_FLOAT(Clamp, result.element_type(), min, operand, max, result);
+}
+
+absl::Status Clamp(const QuantizedTensor& min, const QuantizedTensor& operand,
+                   const QuantizedTensor& max, QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(Clamp, result.storage_type(), result.expressed_type(), min,
+                     operand, max, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/compare.cc

@@ -0,0 +1,164 @@
+/* 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 <type_traits>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+bool CompareHelper(ComparisonDirection comparison_direction, Value lhs,
+                   Value rhs) {
+  switch (comparison_direction) {
+    case ComparisonDirection::kEQ:
+      return (lhs == rhs);
+    case ComparisonDirection::kNE:
+      return (lhs != rhs);
+    case ComparisonDirection::kGE:
+      return (lhs >= rhs);
+    case ComparisonDirection::kGT:
+      return (lhs > rhs);
+    case ComparisonDirection::kLE:
+      return (lhs <= rhs);
+    case ComparisonDirection::kLT:
+      return (lhs < rhs);
+  }
+}
+
+template <typename Value>
+absl::Status CheckParameters(const Value& lhs, const Value& rhs,
+                             ComparisonDirection comparison_direction,
+                             CompareType compare_type, Tensor& result) {
+  if (!(lhs.baseline_element_type() == rhs.baseline_element_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_element_type(lhs) = "
+        "baseline_element_type(rhs)");
+  } else if (!(lhs.shape() == rhs.shape() and lhs.shape() == result.shape())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: shape(lhs) = shape(rhs) = shape(rhs)");
+  } else if (result.element_type() != ElementType::kI1) {
+    return absl::InvalidArgumentError("Expected boolean tensor as result");
+  } else {
+    auto element_type = lhs.element_type();
+    if (!((compare_type == CompareType::kSigned and
+           IsSignedInteger(element_type)) or
+          (compare_type == CompareType::kUnsigned and
+           (IsUnsignedInteger(element_type) or IsBoolean(element_type))) or
+          ((compare_type == CompareType::kFloat or
+            compare_type == CompareType::kTotalOrder) and
+           IsFloat(element_type)))) {
+      return absl::InvalidArgumentError(
+          "Inconsistent compare type vs element type");
+    }
+  }
+  if (compare_type == CompareType::kTotalOrder) {
+    return absl::InvalidArgumentError(
+        "CompareType::kTotalOrder is unsupported");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!(lhs.is_per_tensor_quantized() and rhs.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError("Expected per-tensor quantization");
+    }
+  }
+
+  if (lhs.layout().has_strides() || rhs.layout().has_strides() ||
+      result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status Compare(const Value& lhs, const Value& rhs,
+                     ComparisonDirection comparison_direction,
+                     CompareType compare_type, Tensor& result) {
+  if (auto check =
+          CheckParameters(lhs, rhs, comparison_direction, compare_type, result);
+      !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+  const size_t n = result.num_elements();
+  auto lhs_buffer = lhs.buffer();
+  auto rhs_buffer = rhs.buffer();
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != lhs.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    for (size_t i = 0; i < n; ++i) {
+      auto lhs_value = S::Get(lhs_buffer, i);
+      auto rhs_value = S::Get(rhs_buffer, i);
+      bool result_value =
+          CompareHelper(comparison_direction, lhs_value, rhs_value);
+      Storage<ElementType::kI1>::Set(result_buffer, i, result_value);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    const QuantizedParameter& lhs_quant_param =
+        lhs.type().element_type().parameters(0);
+    const QuantizedParameter& rhs_quant_param =
+        rhs.type().element_type().parameters(0);
+
+    for (size_t i = 0; i < n; ++i) {
+      auto lhs_storage = S::Get(lhs_buffer, i);
+      auto lhs_expressed = Dequantize<storage_type, expressed_type>(
+          lhs_storage, lhs_quant_param);
+
+      auto rhs_storage = S::Get(rhs_buffer, i);
+      auto rhs_expressed = Dequantize<storage_type, expressed_type>(
+          rhs_storage, rhs_quant_param);
+
+      bool result_value =
+          CompareHelper(comparison_direction, lhs_expressed, rhs_expressed);
+      Storage<ElementType::kI1>::Set(result_buffer, i, result_value);
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status Compare(const Tensor& lhs, const Tensor& rhs,
+                     ComparisonDirection comparison_direction,
+                     CompareType compare_type, Tensor& result) {
+  DISPATCH_BOOL_INT_FLOAT(Compare, lhs.element_type(), lhs, rhs,
+                          comparison_direction, compare_type, result);
+}
+
+absl::Status Compare(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     ComparisonDirection comparison_direction,
+                     CompareType compare_type, Tensor& result) {
+  DISPATCH_QUANTIZED(Compare, lhs.storage_type(), lhs.expressed_type(), lhs,
+                     rhs, comparison_direction, compare_type, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/concatenate.cc

@@ -0,0 +1,189 @@
+/* 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 <algorithm>
+#include <cstddef>
+#include <type_traits>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(absl::Span<const Value*> inputs,
+                             DimensionSize dimension, Value& result) {
+  for (auto i = 1; i < inputs.size(); ++i) {
+    if (!(inputs[i]->element_type() == inputs[0]->element_type())) {
+      return absl::InvalidArgumentError(
+          "Constraint violation: same(element_type(inputs...))");
+    }
+  }
+
+  for (size_t ax = 0; ax < inputs[0]->rank(); ++ax) {
+    if (ax != dimension) {
+      for (auto i = 1; i < inputs.size(); ++i) {
+        if (!(inputs[i]->dim(ax) == inputs[0]->dim(ax))) {
+          return absl::InvalidArgumentError(
+              "Constraint violation: same(shape(inputs...)) except for "
+              "dim(inputs..., dimension)");
+        }
+      }
+    }
+  }
+
+  if (inputs.empty()) {
+    return absl::InvalidArgumentError("Constraint violation: 0 < size(inputs)");
+  } else if (!(dimension >= 0 && dimension < inputs[0]->rank())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: 0 <= dimension < rank(inputs[0])");
+  } else if (!(result.element_type() == inputs[0]->element_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: element_type(result) = element_type(inputs[0])");
+  } else {
+    for (size_t ax = 0; ax < result.rank(); ++ax) {
+      DimensionSize expected_dim_size = 0;
+      if (ax == dimension) {
+        for (auto* x : inputs) {
+          expected_dim_size += x->dim(ax);
+        }
+      } else {
+        expected_dim_size = inputs[0]->dim(ax);
+      }
+      if (!(result.dim(ax) == expected_dim_size)) {
+        return absl::InvalidArgumentError(
+            "Constraint violation: element_type(result) = "
+            "element_type(inputs[0])");
+      }
+    }
+  }
+
+  if (std::any_of(inputs.begin(), inputs.end(),
+                  [](auto* x) { return x->layout().has_strides(); }) ||
+      result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status Concatenate(absl::Span<const Value*> inputs,
+                         DimensionSize dimension, Value& result) {
+  if (auto check = CheckParameters(inputs, dimension, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != result.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    DimensionSize dimension_offset = 0;
+    TensorIndex result_index(result.shape());
+    for (auto* input : inputs) {
+      auto input_buffer = input->buffer();
+      for (TensorIndexIterator input_iter{input->shape()};
+           input_iter.has_next(); ++input_iter) {
+        const TensorIndex& input_index = *input_iter;
+        result_index.set(input_index);
+        auto new_dim_size = result_index[dimension] + dimension_offset;
+        result_index.set(dimension, new_dim_size);
+
+        auto linearized_input_index = input_index.linearize();
+        auto linearized_result_index = result_index.linearize();
+
+        auto value = S::Get(input_buffer, linearized_input_index);
+        S::Set(result_buffer, linearized_result_index, value);
+      }
+
+      dimension_offset += input->dim(dimension);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    using ET = typename Storage<expressed_type>::Type;
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    DimensionSize dimension_offset = 0;
+    TensorIndex result_index(result.shape());
+    for (auto* input : inputs) {
+      auto input_buffer = input->buffer();
+      const QuantizedParameter& input_quant_param =
+          input->type().element_type().parameters(0);
+      for (TensorIndexIterator input_iter{input->shape()};
+           input_iter.has_next(); ++input_iter) {
+        const TensorIndex& input_index = *input_iter;
+        result_index.set(input_index);
+        auto new_dim_size = result_index[dimension] + dimension_offset;
+        result_index.set(dimension, new_dim_size);
+
+        auto linearized_input_index = input_index.linearize();
+        auto linearized_result_index = result_index.linearize();
+
+        auto input_storage = S::Get(input_buffer, linearized_input_index);
+        auto result_storage =
+            DequantizeOpQuantizePartial<storage_type, expressed_type>(
+                input_storage, input_quant_param, result_scale_inv,
+                result_quant_param.zero_point, [](auto x) { return x; });
+        S::Set(result_buffer, linearized_result_index, result_storage);
+      }
+
+      dimension_offset += input->dim(dimension);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status Concatenate(absl::Span<const Tensor*> inputs,
+                         DimensionSize dimension, Tensor& result) {
+  DISPATCH_BOOL_INT_FLOAT(Concatenate, result.element_type(), inputs, dimension,
+                          result);
+}
+
+absl::Status Concatenate(absl::Span<const QuantizedTensor*> inputs,
+                         DimensionSize dimension, QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(Concatenate, result.storage_type(),
+                     result.expressed_type(), inputs, dimension, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/debug.cc

@@ -0,0 +1,528 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/legacy/src/debug.h"
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <ostream>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Float>
+inline constexpr Float MaxRelativeDifference() {
+  if constexpr (std::is_same_v<Float, BF16>) {
+    return 1e-2;
+  } else if constexpr (std::is_same_v<Float, F16>) {
+    return 2e-3;
+  } else {
+    return 1e-6;
+  }
+}
+
+// Comare two floats by computing the relative difference. See
+// https://randomascii.wordpress.com/2012/02/25/comparing-floating-point-numbers-2012-edition/
+template <typename Float>
+bool AlmostSame(
+    Float input1, Float input2,
+    Float max_relative_difference = MaxRelativeDifference<Float>()) {
+  bool result;
+  if (std::isnan(input1) || std::isnan(input2)) {
+    result = (std::isnan(input1) and std::isnan(input2));
+  } else if (std::isinf(input1) || std::isinf(input2)) {
+    result = (std::isinf(input1) and std::isinf(input2)) and
+             (std::signbit(static_cast<float>(input1)) ==
+              std::signbit(static_cast<float>(input2)));
+  } else if (input1 == Float(0.0f)) {
+    result = input2 < max_relative_difference;
+  } else if (input2 == Float(0.0f)) {
+    result = input1 < max_relative_difference;
+  } else {
+    Float diff = input1 - input2;
+    if (diff < 0) {
+      diff = -diff;
+    }
+    Float a = (input1 > 0) ? input1 : -input1;
+    Float b = (input2 > 0) ? input2 : -input2;
+    Float largest = (a > b) ? a : b;
+    Float max_diff = (largest * max_relative_difference);
+    result = diff < max_diff;
+  }
+  return result;
+}
+
+template <typename Float>
+bool AlmostSame(
+    const Float* input1, const Float* input2, size_t num_elements,
+    Float max_relative_difference = MaxRelativeDifference<Float>()) {
+  for (size_t i = 0; i < num_elements; ++i) {
+    if (!AlmostSame<Float>(input1[i], input2[i], max_relative_difference)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+}  // namespace
+
+// /////////////////////////////////////////////////////////////////////////////
+
+bool AlmostSame(const Tensor& x, const Tensor& y) {
+  if (x.type() != y.type()) {
+    return false;
+  }
+
+  switch (x.element_type()) {
+    case ElementType::kI1:
+    case ElementType::kSI8:
+    case ElementType::kSI16:
+    case ElementType::kSI32:
+      return !std::memcmp(x.buffer(), y.buffer(), x.num_bytes());
+    case ElementType::kBF16: {
+      using ET = typename Storage<ElementType::kBF16>::Type;
+      return AlmostSame(static_cast<const ET*>(x.buffer()),
+                        static_cast<const ET*>(y.buffer()), x.num_elements());
+    }
+    case ElementType::kF16: {
+      using ET = typename Storage<ElementType::kF16>::Type;
+      return AlmostSame(static_cast<const ET*>(x.buffer()),
+                        static_cast<const ET*>(y.buffer()), x.num_elements());
+    }
+    case ElementType::kF32: {
+      using ET = typename Storage<ElementType::kF32>::Type;
+      return AlmostSame(static_cast<const ET*>(x.buffer()),
+                        static_cast<const ET*>(y.buffer()), x.num_elements());
+    }
+    default:
+      LOG(ERROR) << "Unexpected tensor element type" << x.element_type();
+      return false;
+  }
+}
+
+bool AlmostSame(const QuantizedTensor& x, const QuantizedTensor& y) {
+  if (x.type() != y.type()) {
+    return false;
+  }
+
+  // For now we support only per-tensor quantization.
+  CHECK(x.is_per_tensor_quantized());  // Crash OK
+
+  const QuantizedParameter& quant_param = x.type().element_type().parameters(0);
+
+  auto x_buffer = x.buffer();
+  auto y_buffer = y.buffer();
+
+  size_t n = x.num_elements();
+  for (size_t i = 0; i < n; ++i) {
+    switch (x.storage_type()) {
+      case ElementType::kSI8: {
+        auto x_quantized_value = Storage<ElementType::kSI8>::Get(x_buffer, i);
+        auto y_quantized_value = Storage<ElementType::kSI8>::Get(y_buffer, i);
+        switch (x.expressed_type()) {
+          case ElementType::kBF16: {
+            auto x_value = Dequantize<ElementType::kSI8, ElementType::kBF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI8, ElementType::kBF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF16: {
+            auto x_value = Dequantize<ElementType::kSI8, ElementType::kF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI8, ElementType::kF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF32: {
+            auto x_value = Dequantize<ElementType::kSI8, ElementType::kF32>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI8, ElementType::kF32>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << x.expressed_type();
+            return false;
+        }
+      } break;
+
+      case ElementType::kSI16: {
+        auto x_quantized_value = Storage<ElementType::kSI16>::Get(x_buffer, i);
+        auto y_quantized_value = Storage<ElementType::kSI16>::Get(y_buffer, i);
+        switch (x.expressed_type()) {
+          case ElementType::kBF16: {
+            auto x_value = Dequantize<ElementType::kSI16, ElementType::kBF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI16, ElementType::kBF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF16: {
+            auto x_value = Dequantize<ElementType::kSI16, ElementType::kF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI16, ElementType::kF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF32: {
+            auto x_value = Dequantize<ElementType::kSI16, ElementType::kF32>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI16, ElementType::kF32>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << x.expressed_type();
+            return false;
+        }
+      } break;
+
+      case ElementType::kSI32: {
+        auto x_quantized_value = Storage<ElementType::kSI32>::Get(x_buffer, i);
+        auto y_quantized_value = Storage<ElementType::kSI32>::Get(y_buffer, i);
+        switch (x.expressed_type()) {
+          case ElementType::kBF16: {
+            auto x_value = Dequantize<ElementType::kSI32, ElementType::kBF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI32, ElementType::kBF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF16: {
+            auto x_value = Dequantize<ElementType::kSI32, ElementType::kF16>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI32, ElementType::kF16>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          case ElementType::kF32: {
+            auto x_value = Dequantize<ElementType::kSI32, ElementType::kF32>(
+                x_quantized_value, quant_param);
+            auto y_value = Dequantize<ElementType::kSI32, ElementType::kF32>(
+                y_quantized_value, quant_param);
+            if (!AlmostSame(x_value, y_value)) {
+              return false;
+            }
+          } break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << x.expressed_type();
+            return false;
+        }
+      } break;
+
+      default:
+        LOG(ERROR) << "Unexpected storage type: " << x.storage_type();
+        return false;
+    }
+  }
+
+  return true;
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+std::ostream& operator<<(std::ostream& os, ElementType element_type) {
+  switch (element_type) {
+    case ElementType::kI1:
+      return os << "i1";
+    case ElementType::kSI8:
+      return os << "si8";
+    case ElementType::kSI16:
+      return os << "si16";
+    case ElementType::kSI32:
+      return os << "si32";
+    case ElementType::kBF16:
+      return os << "bf16";
+    case ElementType::kF16:
+      return os << "f16";
+    case ElementType::kF32:
+      return os << "f32";
+    default:
+      LOG(ERROR) << "Unexpected element type: "
+                 << static_cast<int>(element_type);
+      return os;
+  }
+}
+
+std::ostream& operator<<(std::ostream& os, const Shape& s) {
+  for (auto i = 0; i < s.rank(); ++i) {
+    os << s.dim(i) << "x";
+  }
+  return os;
+}
+
+std::ostream& operator<<(std::ostream& os, const TensorType& t) {
+  return os << "tensor<" << t.shape() << t.element_type() << ">";
+}
+
+std::ostream& operator<<(std::ostream& os, const Tensor& t) {
+  os << t.type() << ": ";
+
+  size_t n = t.num_elements();
+  auto t_buffer = t.buffer();
+
+  for (size_t i = 0; i < n; ++i) {
+    if (i > 0) {
+      os << ", ";
+    }
+    switch (t.element_type()) {
+      case ElementType::kI1:
+        os << static_cast<bool>(Storage<ElementType::kI1>::Get(t_buffer, i));
+        break;
+      case ElementType::kSI8:
+        os << static_cast<int>(Storage<ElementType::kSI8>::Get(t_buffer, i));
+        break;
+      case ElementType::kSI16:
+        os << Storage<ElementType::kSI16>::Get(t_buffer, i);
+        break;
+      case ElementType::kSI32:
+        os << Storage<ElementType::kSI32>::Get(t_buffer, i);
+        break;
+      case ElementType::kBF16:
+        os << Storage<ElementType::kBF16>::Get(t_buffer, i);
+        break;
+      case ElementType::kF16:
+        os << Storage<ElementType::kF16>::Get(t_buffer, i);
+        break;
+      case ElementType::kF32:
+        os << Storage<ElementType::kF32>::Get(t_buffer, i);
+        break;
+      default:
+        LOG(ERROR) << "Unexpected element type: " << t.element_type();
+        break;
+    }
+  }
+
+  return os;
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+std::ostream& operator<<(std::ostream& os,
+                         const QuantizedTensorElementType& t) {
+  os << "!quant.uniform<" << t.storage_type();
+  os << ":" << t.expressed_type();
+  if (t.storage_min() || t.storage_max()) {
+    int32_t min;
+    int32_t max;
+    switch (t.storage_type()) {
+      case ElementType::kSI8:
+        min = std::numeric_limits<int8_t>::min();
+        max = std::numeric_limits<int8_t>::max();
+        break;
+      case ElementType::kSI16:
+        min = std::numeric_limits<int16_t>::min();
+        max = std::numeric_limits<int16_t>::max();
+        break;
+      case ElementType::kSI32:
+        min = std::numeric_limits<int32_t>::min();
+        max = std::numeric_limits<int32_t>::max();
+        break;
+      default:
+        LOG(ERROR) << "Unexpected storage type: " << t.storage_type();
+        min = std::numeric_limits<int32_t>::min();
+        max = std::numeric_limits<int32_t>::max();
+        break;
+    }
+    os << "<" << t.storage_min().value_or(min) << ":"
+       << t.storage_max().value_or(max) << ">";
+  }
+  if (t.quantized_dimension()) {
+    os << ":" << *t.quantized_dimension();
+  }
+  os << ", ";
+  if (t.is_per_tensor_quantized()) {
+    const auto& p = t.parameters(0);
+    os << p.scale << ":" << p.zero_point;
+  } else {
+    os << "{";
+    for (size_t i = 0; i < t.num_parameters(); ++i) {
+      if (i > 0) {
+        os << ",";
+      }
+      const auto& p = t.parameters(i);
+      os << p.scale << ":" << p.zero_point;
+    }
+    os << "}";
+  }
+  return os << ">";
+}
+
+std::ostream& operator<<(std::ostream& os, const QuantizedTensorType& t) {
+  return os << "tensor<" << t.shape() << t.element_type() << ">";
+}
+
+std::ostream& operator<<(std::ostream& os, const QuantizedTensor& t) {
+  os << t.type() << ": ";
+
+  if (t.is_per_axis_quantized()) {
+    return os << "<per axis quantization>";
+  }
+
+  const QuantizedParameter& quant_param = t.type().element_type().parameters(0);
+  auto t_buffer = t.buffer();
+
+  size_t n = t.num_elements();
+  for (size_t i = 0; i < n; ++i) {
+    if (i > 0) {
+      os << ", ";
+    }
+
+    switch (t.storage_type()) {
+      case ElementType::kSI8: {
+        auto quantized_value = Storage<ElementType::kSI8>::Get(t_buffer, i);
+        os << static_cast<int>(quantized_value) << "(";
+        switch (t.expressed_type()) {
+          case ElementType::kBF16:
+            os << Dequantize<ElementType::kSI8, ElementType::kBF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF16:
+            os << Dequantize<ElementType::kSI8, ElementType::kF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF32:
+            os << Dequantize<ElementType::kSI8, ElementType::kF32>(
+                quantized_value, quant_param);
+            break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << t.expressed_type();
+        }
+        os << ")";
+      } break;
+
+      case ElementType::kSI16: {
+        auto quantized_value = Storage<ElementType::kSI16>::Get(t_buffer, i);
+        os << quantized_value << "(";
+        switch (t.expressed_type()) {
+          case ElementType::kBF16:
+            os << Dequantize<ElementType::kSI16, ElementType::kBF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF16:
+            os << Dequantize<ElementType::kSI16, ElementType::kF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF32:
+            os << Dequantize<ElementType::kSI16, ElementType::kF32>(
+                quantized_value, quant_param);
+            break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << t.expressed_type();
+        }
+        os << ")";
+      } break;
+
+      case ElementType::kSI32: {
+        auto quantized_value = Storage<ElementType::kSI32>::Get(t_buffer, i);
+        os << quantized_value << "(";
+        switch (t.expressed_type()) {
+          case ElementType::kBF16:
+            os << Dequantize<ElementType::kSI32, ElementType::kBF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF16:
+            os << Dequantize<ElementType::kSI32, ElementType::kF16>(
+                quantized_value, quant_param);
+            break;
+          case ElementType::kF32:
+            os << Dequantize<ElementType::kSI32, ElementType::kF32>(
+                quantized_value, quant_param);
+            break;
+          default:
+            LOG(ERROR) << "Unexpected expressed type: " << t.expressed_type();
+        }
+        os << ")";
+      } break;
+
+      default:
+        LOG(ERROR) << "Unexpected storage type: " << t.storage_type();
+        break;
+    }
+  }
+
+  return os;
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+std::ostream& operator<<(std::ostream& os,
+                         ComparisonDirection comparison_direction) {
+  switch (comparison_direction) {
+    case ComparisonDirection::kEQ:
+      return os << "=";
+    case ComparisonDirection::kNE:
+      return os << "!=";
+    case ComparisonDirection::kGE:
+      return os << ">=";
+    case ComparisonDirection::kGT:
+      return os << ">";
+    case ComparisonDirection::kLE:
+      return os << "<=";
+    case ComparisonDirection::kLT:
+      return os << "<";
+  }
+}
+
+std::ostream& operator<<(std::ostream& os, CompareType compare_type) {
+  switch (compare_type) {
+    case CompareType::kSigned:
+      return os << "signed";
+    case CompareType::kUnsigned:
+      return os << "unsigned";
+    case CompareType::kFloat:
+      return os << "float";
+    case CompareType::kTotalOrder:
+      return os << "totalorder";
+  }
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+std::ostream& operator<<(std::ostream& os, const TensorIndex& tensor_index) {
+  return os << "[" << ToString(tensor_index.index_) << "->"
+            << tensor_index.linear_index_ << "]";
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/debug.h

@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_DEBUG_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_DEBUG_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <ios>
+#include <limits>
+#include <optional>
+#include <ostream>
+#include <sstream>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+bool AlmostSame(const Tensor& x, const Tensor& y);
+bool AlmostSame(const QuantizedTensor& x, const QuantizedTensor& y);
+
+std::ostream& operator<<(std::ostream& os, ElementType element_type);
+std::ostream& operator<<(std::ostream& os, const Shape& shape);
+std::ostream& operator<<(std::ostream& os, const TensorType& tensor_type);
+std::ostream& operator<<(std::ostream& os, const Tensor& tensor);
+std::ostream& operator<<(std::ostream& os, const QuantizedTensorElementType& t);
+std::ostream& operator<<(std::ostream& os,
+                         const QuantizedTensorType& tensor_type);
+std::ostream& operator<<(std::ostream& os, const QuantizedTensor& tensor);
+
+inline std::ostream& operator<<(std::ostream& os, BF16 value) {
+  return os << static_cast<float>(value);
+}
+
+inline std::ostream& operator<<(std::ostream& os, F16 value) {
+  return os << static_cast<float>(value);
+}
+
+std::ostream& operator<<(std::ostream& os,
+                         ComparisonDirection comparison_direction);
+std::ostream& operator<<(std::ostream& os, CompareType compare_type);
+
+std::ostream& operator<<(std::ostream& os, const TensorIndex& tensor_index);
+
+template <typename T>
+std::ostream& operator<<(std::ostream& os, const std::optional<T>& optional) {
+  if (optional) {
+    os << *optional;
+  } else {
+    os << "?";
+  }
+  return os;
+}
+
+template <typename T>
+std::string ToString(
+    const T* ptr, size_t n,
+    size_t max_num_elem_to_print = std::numeric_limits<size_t>::max()) {
+  std::ostringstream os;
+  os << "[";
+  for (size_t i = 0; i < std::min(n, max_num_elem_to_print); ++i) {
+    if (i > 0) {
+      os << ", ";
+    }
+    if constexpr (std::is_same<T, int8_t>::value ||
+                  std::is_same<T, uint8_t>::value) {
+      os << "0x" << std::hex << static_cast<uint32_t>(ptr[i]);
+    } else {
+      os << ptr[i];
+    }
+  }
+  if (n > max_num_elem_to_print) {
+    os << ", ...";
+  }
+  os << "]";
+  return os.str();
+}
+
+template <typename T, typename... Types>
+inline std::string ToString(const std::vector<T, Types...>& vec) {
+  return ToString(vec.data(), vec.size());
+}
+
+template <typename T, typename... Types>
+inline std::ostream& operator<<(std::ostream& os,
+                                const std::vector<T, Types...>& vec) {
+  return os << ToString(vec.data(), vec.size());
+}
+
+template <typename T>
+inline std::string ToString(absl::Span<T> span) {
+  return ToString(span.data(), span.size());
+}
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_DEBUG_H_

tensorflow/lite/experimental/shlo/legacy/src/dispatch.h

@@ -0,0 +1,137 @@
+/* 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_SRC_DISPATCH_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_DISPATCH_H_
+
+namespace stablehlo {
+
+#define DISPATCH_INT(name, element_type, ...)                             \
+  {                                                                       \
+    switch (element_type) {                                               \
+      case ElementType::kSI8:                                             \
+        return name<ElementType::kSI8, ElementType::kSI8>(__VA_ARGS__);   \
+      case ElementType::kSI16:                                            \
+        return name<ElementType::kSI16, ElementType::kSI16>(__VA_ARGS__); \
+      case ElementType::kSI32:                                            \
+        return name<ElementType::kSI32, ElementType::kSI32>(__VA_ARGS__); \
+      default:                                                            \
+        return absl::InvalidArgumentError("Unsupported element type");    \
+    }                                                                     \
+  }
+
+#define DISPATCH_FLOAT(name, element_type, ...)                           \
+  {                                                                       \
+    switch (element_type) {                                               \
+      case ElementType::kBF16:                                            \
+        return name<ElementType::kBF16, ElementType::kBF16>(__VA_ARGS__); \
+      case ElementType::kF16:                                             \
+        return name<ElementType::kF16, ElementType::kF16>(__VA_ARGS__);   \
+      case ElementType::kF32:                                             \
+        return name<ElementType::kF32, ElementType::kF32>(__VA_ARGS__);   \
+      default:                                                            \
+        return absl::InvalidArgumentError("Unsupported element type");    \
+    }                                                                     \
+  }
+
+#define DISPATCH_INT_FLOAT(name, element_type, ...)                       \
+  {                                                                       \
+    switch (element_type) {                                               \
+      case ElementType::kSI8:                                             \
+        return name<ElementType::kSI8, ElementType::kSI8>(__VA_ARGS__);   \
+      case ElementType::kSI16:                                            \
+        return name<ElementType::kSI16, ElementType::kSI16>(__VA_ARGS__); \
+      case ElementType::kSI32:                                            \
+        return name<ElementType::kSI32, ElementType::kSI32>(__VA_ARGS__); \
+      case ElementType::kBF16:                                            \
+        return name<ElementType::kBF16, ElementType::kBF16>(__VA_ARGS__); \
+      case ElementType::kF16:                                             \
+        return name<ElementType::kF16, ElementType::kF16>(__VA_ARGS__);   \
+      case ElementType::kF32:                                             \
+        return name<ElementType::kF32, ElementType::kF32>(__VA_ARGS__);   \
+      default:                                                            \
+        return absl::InvalidArgumentError("Unsupported element type");    \
+    }                                                                     \
+  }
+
+#define DISPATCH_BOOL_INT_FLOAT(name, element_type, ...)                  \
+  {                                                                       \
+    switch (element_type) {                                               \
+      case ElementType::kI1:                                              \
+        return name<ElementType::kI1, ElementType::kI1>(__VA_ARGS__);     \
+      case ElementType::kSI8:                                             \
+        return name<ElementType::kSI8, ElementType::kSI8>(__VA_ARGS__);   \
+      case ElementType::kSI16:                                            \
+        return name<ElementType::kSI16, ElementType::kSI16>(__VA_ARGS__); \
+      case ElementType::kSI32:                                            \
+        return name<ElementType::kSI32, ElementType::kSI32>(__VA_ARGS__); \
+      case ElementType::kBF16:                                            \
+        return name<ElementType::kBF16, ElementType::kBF16>(__VA_ARGS__); \
+      case ElementType::kF16:                                             \
+        return name<ElementType::kF16, ElementType::kF16>(__VA_ARGS__);   \
+      case ElementType::kF32:                                             \
+        return name<ElementType::kF32, ElementType::kF32>(__VA_ARGS__);   \
+      default:                                                            \
+        return absl::InvalidArgumentError("Unsupported element type");    \
+    }                                                                     \
+  }
+
+#define DISPATCH_QUANTIZED(name, storage_type, expressed_type, ...)           \
+  {                                                                           \
+    switch (storage_type) {                                                   \
+      case ElementType::kSI8:                                                 \
+        switch (expressed_type) {                                             \
+          case ElementType::kBF16:                                            \
+            return name<ElementType::kSI8, ElementType::kBF16>(__VA_ARGS__);  \
+          case ElementType::kF16:                                             \
+            return name<ElementType::kSI8, ElementType::kF16>(__VA_ARGS__);   \
+          case ElementType::kF32:                                             \
+            return name<ElementType::kSI8, ElementType::kF32>(__VA_ARGS__);   \
+          default:                                                            \
+            return absl::InvalidArgumentError("Unsupported expressed type");  \
+        }                                                                     \
+        break;                                                                \
+      case ElementType::kSI16:                                                \
+        switch (expressed_type) {                                             \
+          case ElementType::kBF16:                                            \
+            return name<ElementType::kSI16, ElementType::kBF16>(__VA_ARGS__); \
+          case ElementType::kF16:                                             \
+            return name<ElementType::kSI16, ElementType::kF16>(__VA_ARGS__);  \
+          case ElementType::kF32:                                             \
+            return name<ElementType::kSI16, ElementType::kF32>(__VA_ARGS__);  \
+          default:                                                            \
+            return absl::InvalidArgumentError("Unsupported expressed type");  \
+        }                                                                     \
+        break;                                                                \
+      case ElementType::kSI32:                                                \
+        switch (expressed_type) {                                             \
+          case ElementType::kBF16:                                            \
+            return name<ElementType::kSI32, ElementType::kBF16>(__VA_ARGS__); \
+          case ElementType::kF16:                                             \
+            return name<ElementType::kSI32, ElementType::kF16>(__VA_ARGS__);  \
+          case ElementType::kF32:                                             \
+            return name<ElementType::kSI32, ElementType::kF32>(__VA_ARGS__);  \
+          default:                                                            \
+            return absl::InvalidArgumentError("Unsupported expressed type");  \
+        }                                                                     \
+        break;                                                                \
+      default:                                                                \
+        return absl::InvalidArgumentError("Unsupported storage type");        \
+    }                                                                         \
+  }
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_DISPATCH_H_

tensorflow/lite/experimental/shlo/legacy/src/elementwise_binary.cc

@@ -0,0 +1,567 @@
+/* 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 <cstddef>
+#include <type_traits>
+
+#include "absl/status/status.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 {
+
+template <typename Value>
+absl::Status CheckParameters(const Value& lhs, const Value& rhs,
+                             Value& result) {
+  if (!(lhs.baseline_type() == rhs.baseline_type() and
+        lhs.baseline_type() == result.baseline_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_type(on_true) = "
+        "baseline_type(on_false) = baseline_type(result)");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!(lhs.is_per_tensor_quantized() and rhs.is_per_tensor_quantized() and
+          result.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError("Expected per=tensor quantization");
+    }
+  }
+
+  if (lhs.layout().has_strides() || rhs.layout().has_strides() ||
+      result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value,
+          typename Op>
+absl::Status ElementwiseBinaryOp(const Value& lhs, const Value& rhs,
+                                 Value& result, Op&& op) {
+  if (auto check = CheckParameters(lhs, rhs, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  auto lhs_buffer = lhs.buffer();
+  auto rhs_buffer = rhs.buffer();
+  auto result_buffer = result.buffer();
+
+  size_t n = lhs.num_elements();
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != result.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    for (size_t i = 0; i < n; ++i) {
+      auto x = S::Get(lhs_buffer, i);
+      auto y = S::Get(rhs_buffer, i);
+      auto z = op(x, y);
+      S::Set(result_buffer, i, z);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    const QuantizedParameter& lhs_quant_param =
+        lhs.type().element_type().parameters(0);
+    const QuantizedParameter& rhs_quant_param =
+        rhs.type().element_type().parameters(0);
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+
+    using ET = typename Storage<expressed_type>::Type;
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    for (size_t i = 0; i < n; ++i) {
+      auto lhs_storage = S::Get(lhs_buffer, i);
+      auto rhs_storage = S::Get(rhs_buffer, i);
+      auto result_storage =
+          DequantizeOpQuantizePartial<storage_type, expressed_type>(
+              lhs_storage, rhs_storage, lhs_quant_param, rhs_quant_param,
+              result_scale_inv, result_quant_param.zero_point, op);
+      S::Set(result_buffer, i, result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+#define DEFINE_ELEMENTWISE_BINARY_OP(name, element_type, expression)        \
+  absl::Status name(const Tensor& lhs, const Tensor& rhs, Tensor& result) { \
+    return ElementwiseBinaryOp<element_type, element_type>(                 \
+        lhs, rhs, result, [](auto x, auto y) { return expression; });       \
+  }
+
+#define DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(name, storage_type,          \
+                                               expressed_type, expression)  \
+  absl::Status name(const QuantizedTensor& lhs, const QuantizedTensor& rhs, \
+                    QuantizedTensor& result) {                              \
+    return ElementwiseBinaryOp<storage_type, expressed_type>(               \
+        lhs, rhs, result, [](auto x, auto y) { return expression; });       \
+  }
+
+#define DEFINE_ELEMENTWISE_BINARY_OP_BOOL(name, expression) \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_i1, ElementType::kI1, expression);
+
+#define DEFINE_ELEMENTWISE_BINARY_OP_INT(name, expression)                   \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_si8, ElementType::kSI8, expression);   \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_si16, ElementType::kSI16, expression); \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_si32, ElementType::kSI32, expression);
+
+#define DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(name, expression)                   \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_bf16, ElementType::kBF16, expression);   \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_f16, ElementType::kF16, expression);     \
+  DEFINE_ELEMENTWISE_BINARY_OP(name##_f32, ElementType::kF32, expression);     \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(name##_q_si8_bf16, ElementType::kSI8, \
+                                         ElementType::kBF16, expression);      \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(name##_q_si8_f16, ElementType::kSI8,  \
+                                         ElementType::kF16, expression);       \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(name##_q_si8_f32, ElementType::kSI8,  \
+                                         ElementType::kF32, expression);       \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si16_bf16, ElementType::kSI16, ElementType::kBF16, expression); \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si16_f16, ElementType::kSI16, ElementType::kF16, expression);   \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si16_f32, ElementType::kSI16, ElementType::kF32, expression);   \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si32_bf16, ElementType::kSI32, ElementType::kBF16, expression); \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si32_f16, ElementType::kSI32, ElementType::kF16, expression);   \
+  DEFINE_ELEMENTWISE_BINARY_QUANTIZED_OP(                                      \
+      name##_q_si32_f32, ElementType::kSI32, ElementType::kF32, expression);
+
+#define CALL_BINARY_OP_BOOL_HELPER(name, lhs, rhs, result) \
+  case ElementType::kI1:                                   \
+    return name##_i1(lhs, rhs, result);
+
+#define CALL_BINARY_OP_INT_HELPER(name, lhs, rhs, result) \
+  case ElementType::kSI8:                                 \
+    return name##_si8(lhs, rhs, result);                  \
+  case ElementType::kSI16:                                \
+    return name##_si16(lhs, rhs, result);                 \
+  case ElementType::kSI32:                                \
+    return name##_si32(lhs, rhs, result);
+
+#define CALL_BINARY_OP_FLOAT_HELPER(name, lhs, rhs, result) \
+  case ElementType::kBF16:                                  \
+    return name##_bf16(lhs, rhs, result);                   \
+  case ElementType::kF16:                                   \
+    return name##_f16(lhs, rhs, result);                    \
+  case ElementType::kF32:                                   \
+    return name##_f32(lhs, rhs, result);
+
+#define CALL_BINARY_OP_BOOL_INT(name, lhs, rhs, result)                      \
+  {                                                                          \
+    auto element_type = lhs.element_type();                                  \
+    switch (element_type) {                                                  \
+      CALL_BINARY_OP_BOOL_HELPER(name, lhs, rhs, result);                    \
+      CALL_BINARY_OP_INT_HELPER(name, lhs, rhs, result);                     \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_BINARY_OP_INT(name, lhs, rhs, result)                           \
+  {                                                                          \
+    auto element_type = lhs.element_type();                                  \
+    switch (element_type) {                                                  \
+      CALL_BINARY_OP_INT_HELPER(name, lhs, rhs, result);                     \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_BINARY_OP_INT_FLOAT(name, lhs, rhs, result)                     \
+  {                                                                          \
+    auto element_type = lhs.element_type();                                  \
+    switch (element_type) {                                                  \
+      CALL_BINARY_OP_INT_HELPER(name, lhs, rhs, result);                     \
+      CALL_BINARY_OP_FLOAT_HELPER(name, lhs, rhs, result);                   \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_BINARY_OP_FLOAT(name, lhs, rhs, result)                         \
+  {                                                                          \
+    auto element_type = lhs.element_type();                                  \
+    switch (element_type) {                                                  \
+      CALL_BINARY_OP_FLOAT_HELPER(name, lhs, rhs, result);                   \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_BINARY_OP_BOOL_INT_FLOAT(name, lhs, rhs, result)                \
+  {                                                                          \
+    auto element_type = lhs.element_type();                                  \
+    switch (element_type) {                                                  \
+      CALL_BINARY_OP_BOOL_HELPER(name, lhs, rhs, result);                    \
+      CALL_BINARY_OP_INT_HELPER(name, lhs, rhs, result);                     \
+      CALL_BINARY_OP_FLOAT_HELPER(name, lhs, rhs, result);                   \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_BINARY_QUANTIZED_OP(name, lhs, rhs, result)                    \
+  {                                                                         \
+    auto storage_type = lhs.storage_type();                                 \
+    auto expressed_type = lhs.expressed_type();                             \
+    switch (storage_type) {                                                 \
+      case ElementType::kSI8:                                               \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si8_bf16(lhs, rhs, result);                     \
+          case ElementType::kF16:                                           \
+            return name##_q_si8_f16(lhs, rhs, result);                      \
+          case ElementType::kF32:                                           \
+            return name##_q_si8_f32(lhs, rhs, result);                      \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      case ElementType::kSI16:                                              \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si16_bf16(lhs, rhs, result);                    \
+          case ElementType::kF16:                                           \
+            return name##_q_si16_f16(lhs, rhs, result);                     \
+          case ElementType::kF32:                                           \
+            return name##_q_si16_f32(lhs, rhs, result);                     \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      case ElementType::kSI32:                                              \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si32_bf16(lhs, rhs, result);                    \
+          case ElementType::kF16:                                           \
+            return name##_q_si32_f16(lhs, rhs, result);                     \
+          case ElementType::kF32:                                           \
+            return name##_q_si32_f32(lhs, rhs, result);                     \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      default:                                                              \
+        return absl::InvalidArgumentError("Unexpected storage type");       \
+    }                                                                       \
+  }
+
+}  // namespace
+
+// /////////////////////////////////////////////////////////////////////////////
+// Add
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Add, x or y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Add, x + y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Add, x + y);
+
+}  // namespace
+
+absl::Status Add(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT_FLOAT(Add, lhs, rhs, result);
+}
+
+absl::Status Add(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                 QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Add, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// And
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(And, x&& y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(And, x& y);
+
+}  // namespace
+
+absl::Status And(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT(And, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Atan2
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Atan2, std::atan2(static_cast<float>(x),
+                                                     static_cast<float>(y)));
+
+}  // namespace
+
+absl::Status Atan2(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_FLOAT(Atan2, lhs, rhs, result);
+}
+
+absl::Status Atan2(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                   QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Atan2, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Divide
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Divide, x / y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Divide, x / y);
+
+}  // namespace
+
+absl::Status Divide(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_INT_FLOAT(Divide, lhs, rhs, result);
+}
+
+absl::Status Divide(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                    QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Divide, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Maximum
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Maximum, x or y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Maximum, (x > y) ? x : y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Maximum, (x > y) ? x : y);
+
+}  // namespace
+
+absl::Status Maximum(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT_FLOAT(Maximum, lhs, rhs, result);
+}
+
+absl::Status Maximum(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Maximum, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Minimum
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Minimum, x and y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Minimum, (x > y) ? y : x);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Minimum, (x > y) ? y : x);
+
+}  // namespace
+
+absl::Status Minimum(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT_FLOAT(Minimum, lhs, rhs, result);
+}
+
+absl::Status Minimum(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                     QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Minimum, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Multiply
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Multiply, x and y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Multiply, x* y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Multiply, x* y);
+
+}  // namespace
+
+absl::Status Multiply(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT_FLOAT(Multiply, lhs, rhs, result);
+}
+
+absl::Status Multiply(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                      QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Multiply, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Or
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Or, x or y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Or, x | y);
+
+}  // namespace
+
+absl::Status Or(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT(Or, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Power
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Power, std::pow(static_cast<float>(x),
+                                                 static_cast<int>(y)));
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Power, std::powf(static_cast<float>(x),
+                                                    static_cast<float>(y)));
+
+}  // namespace
+
+absl::Status Power(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_INT_FLOAT(Power, lhs, rhs, result);
+}
+
+absl::Status Power(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                   QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Power, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Remainder
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Remainder, x % y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Remainder, std::fmod(static_cast<float>(x),
+                                                        static_cast<float>(y)));
+
+}  // namespace
+
+absl::Status Remainder(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_INT_FLOAT(Remainder, lhs, rhs, result);
+}
+
+absl::Status Remainder(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                       QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Remainder, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// ShiftLeft
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(ShiftLeft, x << y);
+
+}  // namespace
+
+absl::Status ShiftLeft(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_INT(ShiftLeft, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// ShiftRightArithmetic
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(ShiftRightArithmetic, x >> y);
+
+}  // namespace
+
+absl::Status ShiftRightArithmetic(const Tensor& lhs, const Tensor& rhs,
+                                  Tensor& result) {
+  CALL_BINARY_OP_INT(ShiftRightArithmetic, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// ShiftRightLogical
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Int>
+inline Int ShiftRightLogical(Int x, Int y) {
+  using UInt = typename std::make_unsigned<Int>::type;
+  return static_cast<UInt>(x) >> y;
+}
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(ShiftRightLogical, ShiftRightLogical(x, y));
+
+}  // namespace
+
+absl::Status ShiftRightLogical(const Tensor& lhs, const Tensor& rhs,
+                               Tensor& result) {
+  CALL_BINARY_OP_INT(ShiftRightLogical, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Subtract
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Subtract, x - y);
+DEFINE_ELEMENTWISE_BINARY_OP_FLOAT(Subtract, x - y);
+
+}  // namespace
+
+absl::Status Subtract(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_INT_FLOAT(Subtract, lhs, rhs, result);
+}
+
+absl::Status Subtract(const QuantizedTensor& lhs, const QuantizedTensor& rhs,
+                      QuantizedTensor& result) {
+  CALL_BINARY_QUANTIZED_OP(Subtract, lhs, rhs, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Xor
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_BINARY_OP_BOOL(Xor, x xor y);
+DEFINE_ELEMENTWISE_BINARY_OP_INT(Xor, x ^ y);
+
+}  // namespace
+
+absl::Status Xor(const Tensor& lhs, const Tensor& rhs, Tensor& result) {
+  CALL_BINARY_OP_BOOL_INT(Xor, lhs, rhs, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/elementwise_unary.cc

@@ -0,0 +1,749 @@
+/* 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 <bit>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+#include <version>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(const Value& operand, Value& result) {
+  if (operand.baseline_type() != result.baseline_type()) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_type(operand) = baseline_type(result)");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!(operand.is_per_tensor_quantized() and
+          result.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError("Expected per-tensor quantization");
+    }
+  }
+
+  if (operand.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value,
+          typename Op>
+absl::Status ElementwiseUnaryOp(const Value& operand, Value& result, Op&& op) {
+  if (auto check = CheckParameters(operand, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  auto operand_buffer = operand.buffer();
+  auto result_buffer = result.buffer();
+
+  size_t n = operand.num_elements();
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != operand.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    for (size_t i = 0; i < n; ++i) {
+      auto x = S::Get(operand_buffer, i);
+      auto y = op(x);
+      S::Set(result_buffer, i, y);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    const QuantizedParameter& operand_quant_param =
+        operand.type().element_type().parameters(0);
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+
+    using ET = typename Storage<expressed_type>::Type;
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    for (size_t i = 0; i < n; ++i) {
+      auto operand_storage = S::Get(operand_buffer, i);
+      auto result_storage =
+          DequantizeOpQuantizePartial<storage_type, expressed_type>(
+              operand_storage, operand_quant_param, result_scale_inv,
+              result_quant_param.zero_point, op);
+      S::Set(result_buffer, i, result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+#define DEFINE_ELEMENTWISE_UNARY_OP(name, element_type, expression) \
+  absl::Status name(const Tensor& operand, Tensor& result) {        \
+    return ElementwiseUnaryOp<element_type, element_type>(          \
+        operand, result, [](auto x) { return expression; });        \
+  }
+
+#define DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name, storage_type,              \
+                                              expressed_type, expression)      \
+  absl::Status name(const QuantizedTensor& operand, QuantizedTensor& result) { \
+    return ElementwiseUnaryOp<storage_type, expressed_type>(                   \
+        operand, result, [](auto x) { return expression; });                   \
+  }
+
+#define DEFINE_ELEMENTWISE_UNARY_OP_BOOL(name, expression) \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_i1, ElementType::kI1, expression);
+
+#define DEFINE_ELEMENTWISE_UNARY_OP_INT(name, expression)                   \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_si8, ElementType::kSI8, expression);   \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_si16, ElementType::kSI16, expression); \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_si32, ElementType::kSI32, expression);
+
+#define DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(name, expression)                    \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_bf16, ElementType::kBF16, expression);    \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_f16, ElementType::kF16, expression);      \
+  DEFINE_ELEMENTWISE_UNARY_OP(name##_f32, ElementType::kF32, expression);      \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si8_bf16, ElementType::kSI8,  \
+                                        ElementType::kBF16, expression);       \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si8_f16, ElementType::kSI8,   \
+                                        ElementType::kF16, expression);        \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si8_f32, ElementType::kSI8,   \
+                                        ElementType::kF32, expression);        \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(                                       \
+      name##_q_si16_bf16, ElementType::kSI16, ElementType::kBF16, expression); \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si16_f16, ElementType::kSI16, \
+                                        ElementType::kF16, expression);        \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si16_f32, ElementType::kSI16, \
+                                        ElementType::kF32, expression);        \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(                                       \
+      name##_q_si32_bf16, ElementType::kSI32, ElementType::kBF16, expression); \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si32_f16, ElementType::kSI32, \
+                                        ElementType::kF16, expression);        \
+  DEFINE_ELEMENTWISE_UNARY_QUANTIZED_OP(name##_q_si32_f32, ElementType::kSI32, \
+                                        ElementType::kF32, expression);
+
+#define CALL_UNARY_OP_BOOL_HELPER(name, operand, result) \
+  case ElementType::kI1:                                 \
+    return name##_i1(operand, result);
+
+#define CALL_UNARY_OP_INT_HELPER(name, operand, result) \
+  case ElementType::kSI8:                               \
+    return name##_si8(operand, result);                 \
+  case ElementType::kSI16:                              \
+    return name##_si16(operand, result);                \
+  case ElementType::kSI32:                              \
+    return name##_si32(operand, result);
+
+#define CALL_UNARY_OP_FLOAT_HELPER(name, operand, result) \
+  case ElementType::kBF16:                                \
+    return name##_bf16(operand, result);                  \
+  case ElementType::kF16:                                 \
+    return name##_f16(operand, result);                   \
+  case ElementType::kF32:                                 \
+    return name##_f32(operand, result);
+
+#define CALL_UNARY_OP_BOOL_INT(name, operand, result)                        \
+  {                                                                          \
+    auto element_type = operand.element_type();                              \
+    switch (element_type) {                                                  \
+      CALL_UNARY_OP_BOOL_HELPER(name, operand, result);                      \
+      CALL_UNARY_OP_INT_HELPER(name, operand, result);                       \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_UNARY_OP_INT(name, operand, result)                             \
+  {                                                                          \
+    auto element_type = operand.element_type();                              \
+    switch (element_type) {                                                  \
+      CALL_UNARY_OP_INT_HELPER(name, operand, result);                       \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_UNARY_OP_FLOAT(name, operand, result)                           \
+  {                                                                          \
+    auto element_type = operand.element_type();                              \
+    switch (element_type) {                                                  \
+      CALL_UNARY_OP_FLOAT_HELPER(name, operand, result);                     \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_UNARY_OP_INT_FLOAT(name, operand, result)                       \
+  {                                                                          \
+    auto element_type = operand.element_type();                              \
+    switch (element_type) {                                                  \
+      CALL_UNARY_OP_INT_HELPER(name, operand, result);                       \
+      CALL_UNARY_OP_FLOAT_HELPER(name, operand, result);                     \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_UNARY_OP_BOOL_INT_FLOAT(name, operand, result)                  \
+  {                                                                          \
+    auto element_type = operand.element_type();                              \
+    switch (element_type) {                                                  \
+      CALL_UNARY_OP_BOOL_HELPER(name, operand, result);                      \
+      CALL_UNARY_OP_INT_HELPER(name, operand, result);                       \
+      CALL_UNARY_OP_FLOAT_HELPER(name, operand, result);                     \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unexpected tensor element type"); \
+    }                                                                        \
+  }
+
+#define CALL_UNARY_QUANTIZED_OP(name, operand, result)                      \
+  {                                                                         \
+    auto storage_type = operand.storage_type();                             \
+    auto expressed_type = operand.expressed_type();                         \
+    switch (storage_type) {                                                 \
+      case ElementType::kSI8:                                               \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si8_bf16(operand, result);                      \
+          case ElementType::kF16:                                           \
+            return name##_q_si8_f16(operand, result);                       \
+          case ElementType::kF32:                                           \
+            return name##_q_si8_f32(operand, result);                       \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      case ElementType::kSI16:                                              \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si16_bf16(operand, result);                     \
+          case ElementType::kF16:                                           \
+            return name##_q_si16_f16(operand, result);                      \
+          case ElementType::kF32:                                           \
+            return name##_q_si16_f32(operand, result);                      \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      case ElementType::kSI32:                                              \
+        switch (expressed_type) {                                           \
+          case ElementType::kBF16:                                          \
+            return name##_q_si32_bf16(operand, result);                     \
+          case ElementType::kF16:                                           \
+            return name##_q_si32_f16(operand, result);                      \
+          case ElementType::kF32:                                           \
+            return name##_q_si32_f32(operand, result);                      \
+          default:                                                          \
+            return absl::InvalidArgumentError("Unexpected expressed type"); \
+        }                                                                   \
+      default:                                                              \
+        return absl::InvalidArgumentError("Unexpected storage type");       \
+    }                                                                       \
+  }
+
+}  // namespace
+
+// /////////////////////////////////////////////////////////////////////////////
+// Abs
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_UNARY_OP_INT(Abs, ((x > 0) ? x : -x));
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Abs, ((x > 0) ? x : -x));
+
+}  // namespace
+
+absl::Status Abs(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_INT_FLOAT(Abs, operand, result);
+}
+
+absl::Status Abs(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Abs, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Cbrt
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Cbrt, std::cbrt(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Cbrt(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Cbrt, operand, result);
+}
+
+absl::Status Cbrt(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Cbrt, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Ceil
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Ceil, std::ceil(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Ceil(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Ceil, operand, result);
+}
+
+absl::Status Ceil(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Ceil, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Cosine
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Cosine, std::cos(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Cosine(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Cosine, operand, result);
+}
+
+absl::Status Cosine(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Cosine, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// CountLeadingZeros
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Int>
+inline Int CountLeadingZeros(Int x) {
+  using UInt = typename std::make_unsigned<Int>::type;
+#if __cpp_lib_bitops >= 201907L
+  return std::countl_zero(static_cast<UInt>(x));
+#else
+  if (!x) {
+    return 8 * sizeof(x);
+  }
+  Int result = 0;
+  auto mask = UInt(1) << (8 * (sizeof(x) - 1) + 7);
+  for (auto t = static_cast<UInt>(x); t > 0; t <<= 1) {
+    if (t & mask) break;
+    result++;
+  }
+  return result;
+#endif
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_INT(CountLeadingZeros, CountLeadingZeros(x));
+
+}  // namespace
+
+absl::Status CountLeadingZeros(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_INT(CountLeadingZeros, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Exponential
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Exponential, std::exp(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Exponential(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Exponential, operand, result);
+}
+
+absl::Status Exponential(const QuantizedTensor& operand,
+                         QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Exponential, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// ExponentialMinusOne
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(ExponentialMinusOne,
+                                  std::expm1(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status ExponentialMinusOne(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(ExponentialMinusOne, operand, result);
+}
+
+absl::Status ExponentialMinusOne(const QuantizedTensor& operand,
+                                 QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(ExponentialMinusOne, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Floor
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Floor, std::floor(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Floor(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Floor, operand, result);
+}
+
+absl::Status Floor(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Floor, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Log
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Log, std::log(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Log(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Log, operand, result);
+}
+
+absl::Status Log(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Log, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// LogPlusOne
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(LogPlusOne,
+                                  std::log1p(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status LogPlusOne(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(LogPlusOne, operand, result);
+}
+
+absl::Status LogPlusOne(const QuantizedTensor& operand,
+                        QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(LogPlusOne, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Logistic
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Logistic,
+                                  1.0f / (1.0f +
+                                          std::exp(static_cast<float>(-x))));
+
+}  // namespace
+
+absl::Status Logistic(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Logistic, operand, result);
+}
+
+absl::Status Logistic(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Logistic, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Negate
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_UNARY_OP_INT(Negate, -x);
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Negate, -x);
+
+}  // namespace
+
+absl::Status Negate(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_INT_FLOAT(Negate, operand, result);
+}
+
+absl::Status Negate(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Negate, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Not
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+DEFINE_ELEMENTWISE_UNARY_OP_BOOL(Not, !x);
+DEFINE_ELEMENTWISE_UNARY_OP_INT(Not, ~x);
+
+}  // namespace
+
+absl::Status Not(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_BOOL_INT(Not, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Popcnt
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Int>
+Int Popcount(Int x) {
+#if __cpp_lib_bitops >= 201907L
+  return std::popcount(static_cast<uint32_t>(x));
+#else
+  using UInt = typename std::make_unsigned<Int>::type;
+  Int result = 0;
+  UInt mask = 0x1;
+  for (auto t = static_cast<UInt>(x); t > 0; t >>= 1) {
+    result += (t & mask);
+  }
+  return result;
+#endif
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_INT(Popcnt, Popcount(x));
+
+}  // namespace
+
+absl::Status Popcnt(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_INT(Popcnt, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// RoundNearestAfz
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Float>
+inline Float RoundNearestAfz(Float x) {
+  // TODO(cbasile): Performing the op with a conversion to float is
+  // inefficient for bf16 and f16 types.
+  return std::round(static_cast<float>(x));
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(RoundNearestAfz, RoundNearestAfz(x));
+
+}  // namespace
+
+absl::Status RoundNearestAfz(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(RoundNearestAfz, operand, result);
+}
+
+absl::Status RoundNearestAfz(const QuantizedTensor& operand,
+                             QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(RoundNearestAfz, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// RoundNearestEven
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Float>
+inline Float RoundNearestEven(Float x) {
+  // TODO(cbasile): Performing the op with a conversion to float is
+  // inefficient for bf16 and f16 types.
+  return x - static_cast<Float>(std::remainder(static_cast<float>(x), 1.0f));
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(RoundNearestEven, RoundNearestEven(x));
+
+}  // namespace
+
+absl::Status RoundNearestEven(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(RoundNearestEven, operand, result);
+}
+
+absl::Status RoundNearestEven(const QuantizedTensor& operand,
+                              QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(RoundNearestEven, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Rsqrt
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Float>
+inline Float Rsqrt(Float x) {
+  // TODO(cbasile): Performing the op with a conversion to float is
+  // inefficient for bf16 and f16 types.
+  return Float{1} / static_cast<Float>(std::sqrt(static_cast<float>(x)));
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Rsqrt, Rsqrt(x));
+
+}  // namespace
+
+absl::Status Rsqrt(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Rsqrt, operand, result);
+}
+
+absl::Status Rsqrt(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Rsqrt, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Sign
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+template <typename Number>
+inline Number Sign(Number x) {
+  if constexpr (std::is_integral<Number>::value) {
+    return x < 0 ? -1 : (x > 0 ? 1 : 0);
+  } else {
+    static_assert(std::is_floating_point<Number>::value ||
+                  std::is_same_v<Number, BF16> || std::is_same_v<Number, F16>);
+    if (std::isnan(x)) {
+      return NAN;
+    }
+    return (x < 0 ? -1 : (x > 0 ? 1 : 0));
+  }
+}
+
+DEFINE_ELEMENTWISE_UNARY_OP_INT(Sign, Sign(x));
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Sign, Sign(x));
+
+}  // namespace
+
+absl::Status Sign(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_INT_FLOAT(Sign, operand, result);
+}
+
+absl::Status Sign(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Sign, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Sine
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is
+// inefficient for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Sine, std::sin(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Sine(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Sine, operand, result);
+}
+
+absl::Status Sine(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Sine, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Sqrt
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is inefficient
+// for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Sqrt, std::sqrt(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Sqrt(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Sqrt, operand, result);
+}
+
+absl::Status Sqrt(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Sqrt, operand, result);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+// Tanh
+// /////////////////////////////////////////////////////////////////////////////
+
+namespace {
+
+// TODO(cbasile): Performing the op with a conversion to float is inefficient
+// for bf16 and f16 types.
+DEFINE_ELEMENTWISE_UNARY_OP_FLOAT(Tanh, std::tanh(static_cast<float>(x)));
+
+}  // namespace
+
+absl::Status Tanh(const Tensor& operand, Tensor& result) {
+  CALL_UNARY_OP_FLOAT(Tanh, operand, result);
+}
+
+absl::Status Tanh(const QuantizedTensor& operand, QuantizedTensor& result) {
+  CALL_UNARY_QUANTIZED_OP(Tanh, operand, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/f16.h

@@ -0,0 +1,36 @@
+/* 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_SRC_F16_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_F16_H_
+
+#include "tensorflow/lite/experimental/shlo/legacy/src/has_keyword.h"
+
+#if defined(__STDCPP_FLOAT16_T__)
+#include <stdfloat>
+namespace stablehlo {
+using F16 = float16_t;
+}  // namespace stablehlo
+
+#elif __has_keyword(_Float16)
+namespace stablehlo {
+using F16 = _Float16;
+}  // namespace stablehlo
+
+#else
+#error Type F16 is not available
+#endif
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_F16_H_

tensorflow/lite/experimental/shlo/legacy/src/has_keyword.h

@@ -0,0 +1,32 @@
+/* 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_SRC_HAS_KEYWORD_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_HAS_KEYWORD_H_
+
+// CAUTION: __is_identifier behaves opposite how you would expect!
+// '__is_identifier' returns '0' if '__x' is a reserved identifier provided by
+// the compiler and '1' otherwise.  borrowed from LLVM __config header under
+// Apache license 2.
+// (https://www.mend.io/blog/top-10-apache-license-questions-answered/)
+
+#ifndef __is_identifier       // Optional of course.
+#define __is_identifier(x) 1  // Compatibility with non-clang compilers.
+#endif
+
+// More sensible macro for keyword detection
+#define __has_keyword(__x) !(__is_identifier(__x))
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_HAS_KEYWORD_H_

tensorflow/lite/experimental/shlo/legacy/src/iota.cc

@@ -0,0 +1,99 @@
+/* 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 <type_traits>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(DimensionSize iota_dimension,
+                             const Value& result) {
+  if (!(0 <= iota_dimension && iota_dimension < result.rank())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: 0 <= iota_dimension < rank(result)");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status Iota(DimensionSize iota_dimension, Value& result) {
+  if (auto check = CheckParameters(iota_dimension, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+  using ST = typename S::Type;
+
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != result.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    for (TensorIndexIterator iter(result.shape()); iter.has_next(); ++iter) {
+      auto& index = *iter;
+      ST y = index[iota_dimension];
+      S::Set(result_buffer, index.linearize(), y);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+
+    using ET = typename Storage<expressed_type>::Type;
+
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    for (TensorIndexIterator iter(result.shape()); iter.has_next(); ++iter) {
+      auto& index = *iter;
+      ET result_expressed = index[iota_dimension];
+      auto result_storage = QuantizePartial<storage_type, expressed_type>(
+          result_expressed, result_scale_inv, result_quant_param.zero_point);
+      S::Set(result_buffer, index.linearize(), result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status Iota(DimensionSize iota_dimension, Tensor& result) {
+  DISPATCH_INT_FLOAT(Iota, result.element_type(), iota_dimension, result);
+}
+
+absl::Status Iota(DimensionSize iota_dimension, QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(Iota, result.storage_type(), result.expressed_type(),
+                     iota_dimension, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/is_finite.cc

@@ -0,0 +1,104 @@
+/* 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 <cstddef>
+#include <type_traits>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(const Value& operand, Tensor& result) {
+  if (operand.shape() != result.shape()) {
+    return absl::InvalidArgumentError("Inconsistent input/output shapes");
+  } else if (result.element_type() != ElementType::kI1) {
+    return absl::InvalidArgumentError("Unexpected output tensor element type");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!operand.is_per_tensor_quantized()) {
+      return absl::InvalidArgumentError("Expected per-tensor quantization");
+    }
+  }
+
+  if (operand.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status IsFinite(const Value& operand, Tensor& result) {
+  if (auto check = CheckParameters(operand, result); !check.ok()) {
+    return check;
+  }
+
+  using S = Storage<storage_type>;
+
+  auto operand_buffer = operand.buffer();
+  auto result_buffer = result.buffer();
+
+  size_t n = operand.num_elements();
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != operand.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    for (size_t i = 0; i < n; ++i) {
+      auto x = S::Get(operand_buffer, i);
+      auto y = std::isfinite(static_cast<float>(x));
+      Storage<ElementType::kI1>::Set(result_buffer, i, y);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    const QuantizedParameter& operand_quant_param =
+        operand.type().element_type().parameters(0);
+
+    for (size_t i = 0; i < n; ++i) {
+      auto operand_storage = S::Get(operand_buffer, i);
+      auto operand_expressed = Dequantize<storage_type, expressed_type>(
+          operand_storage, operand_quant_param);
+      auto result_expressed =
+          std::isfinite(static_cast<float>(operand_expressed));
+      Storage<ElementType::kI1>::Set(result_buffer, i, result_expressed);
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status IsFinite(const Tensor& operand, Tensor& result) {
+  DISPATCH_FLOAT(IsFinite, operand.element_type(), operand, result);
+}
+
+absl::Status IsFinite(const QuantizedTensor& operand, Tensor& result) {
+  DISPATCH_QUANTIZED(IsFinite, operand.storage_type(), operand.expressed_type(),
+                     operand, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/select.cc

@@ -0,0 +1,161 @@
+/* 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 <type_traits>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+template <typename Value>
+absl::Status CheckParameters(const Tensor& pred, const Value& on_true,
+                             const Value& on_false, Value& result) {
+  if (!(pred.rank() == 0 or pred.shape() == on_true.shape())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: rank(pred) = 0 or shape(pred) = "
+        "shape(on_true)");
+  } else if (!(on_true.baseline_type() == on_false.baseline_type() and
+               on_true.baseline_type() == result.baseline_type())) {
+    return absl::InvalidArgumentError(
+        "Constraint violation: baseline_type(on_true) = "
+        "baseline_type(on_false) = baseline_type(result)");
+  } else if (pred.element_type() != ElementType::kI1) {
+    return absl::InvalidArgumentError("Expected boolean tensor as predicate");
+  }
+
+  if constexpr (std::is_same_v<Value, QuantizedTensor>) {
+    if (!(on_true.is_per_tensor_quantized() and
+          on_false.is_per_tensor_quantized() and
+          result.is_per_tensor_quantized())) {
+      return absl::InvalidArgumentError("Expected per-tensor quantization");
+    }
+  }
+
+  if (pred.layout().has_strides() || on_true.layout().has_strides() ||
+      on_false.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Value>
+absl::Status Select(const Tensor& pred, const Value& on_true,
+                    const Value& on_false, Value& result) {
+  if (auto check = CheckParameters(pred, on_true, on_false, result);
+      !check.ok()) {
+    return check;
+  }
+
+  using P = Storage<ElementType::kI1>;
+  using S = Storage<storage_type>;
+
+  const bool pred_is_tensor = (pred.rank() > 0);
+  const size_t n = result.num_elements();
+
+  auto pred_buffer = pred.buffer();
+  auto on_true_buffer = on_true.buffer();
+  auto on_false_buffer = on_false.buffer();
+  auto result_buffer = result.buffer();
+
+  if constexpr (std::is_same_v<Value, Tensor>) {
+    if (storage_type != result.element_type()) {
+      return absl::InvalidArgumentError("Unexpected tensor element type");
+    }
+
+    bool selection_value;
+    for (size_t i = 0; i < n; ++i) {
+      if (pred_is_tensor || (i == 0)) {
+        selection_value = P::Get(pred_buffer, i);
+      }
+      auto input_buffer = selection_value ? on_true_buffer : on_false_buffer;
+      auto result_value = S::Get(input_buffer, i);
+      S::Set(result_buffer, i, result_value);
+    }
+
+  } else {
+    static_assert(std::is_same_v<Value, QuantizedTensor>);
+
+    if (storage_type != result.storage_type()) {
+      return absl::InvalidArgumentError("Unexpected storage type");
+    } else if (expressed_type != result.expressed_type()) {
+      return absl::InvalidArgumentError("Unexpected expressed type");
+    }
+
+    using ET = typename Storage<expressed_type>::Type;
+
+    const QuantizedParameter& on_true_quant_param =
+        on_true.type().element_type().parameters(0);
+    const QuantizedParameter& on_false_quant_param =
+        on_false.type().element_type().parameters(0);
+    const QuantizedParameter& result_quant_param =
+        result.type().element_type().parameters(0);
+    ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+    bool selection_value;
+    for (size_t i = 0; i < n; ++i) {
+      if (pred_is_tensor || (i == 0)) {
+        selection_value = P::Get(pred_buffer, i);
+      }
+
+      const void* input_buffer;
+      const QuantizedParameter* input_quant_param;
+      if (selection_value) {
+        input_buffer = on_true_buffer;
+        input_quant_param = &on_true_quant_param;
+      } else {
+        input_buffer = on_false_buffer;
+        input_quant_param = &on_false_quant_param;
+      }
+
+      auto input_storage = S::Get(input_buffer, i);
+      auto result_storage =
+          DequantizeOpQuantizePartial<storage_type, expressed_type>(
+              input_storage, *input_quant_param, result_scale_inv,
+              result_quant_param.zero_point, [](auto x) { return x; });
+      S::Set(result_buffer, i, result_storage);
+    }
+
+    if (auto status = CompleteQuantization<storage_type>(result);
+        !status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status Select(const Tensor& pred, const Tensor& on_true,
+                    const Tensor& on_false, Tensor& result) {
+  DISPATCH_BOOL_INT_FLOAT(Select, result.element_type(), pred, on_true,
+                          on_false, result);
+}
+
+absl::Status Select(const Tensor& pred, const QuantizedTensor& on_true,
+                    const QuantizedTensor& on_false, QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(Select, result.storage_type(), result.expressed_type(),
+                     pred, on_true, on_false, result);
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/shlo.cc

@@ -0,0 +1,126 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+
+#include <cstddef>
+#include <cstring>
+#include <functional>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+
+namespace stablehlo {
+
+Axes Shape::axes() const {
+  Axes tmp(dims_.size());
+  std::iota(tmp.begin(), tmp.end(), 0);
+  return tmp;
+}
+
+size_t Shape::num_elements() const {
+  if (dims_.empty()) {
+    return 0;
+  }
+  return std::accumulate(dims_.begin(), dims_.end(), DimensionSize{1},
+                         std::multiplies<DimensionSize>());
+}
+
+namespace {
+
+template <ElementType element_type>
+struct GetNumBytes {
+  size_t operator()() {
+    return sizeof(typename stablehlo::Storage<element_type>::Type);
+  }
+};
+
+template <template <ElementType> typename Op, typename Result>
+Result CallTemplatedFunctorWithResult(ElementType element_type) {
+  switch (element_type) {
+    case ElementType::kI1:
+      return Op<ElementType::kI1>()();
+    case ElementType::kSI8:
+      return Op<ElementType::kSI8>()();
+    case ElementType::kSI16:
+      return Op<ElementType::kSI16>()();
+    case ElementType::kSI32:
+      return Op<ElementType::kSI32>()();
+    case ElementType::kBF16:
+      return Op<ElementType::kBF16>()();
+    case ElementType::kF16:
+      return Op<ElementType::kF16>()();
+    case ElementType::kF32:
+      return Op<ElementType::kF32>()();
+    default:
+      LOG(ERROR) << "Unexpected tensor element type: "
+                 << static_cast<int>(element_type);
+      return Result();
+  }
+}
+
+}  // namespace
+
+size_t TensorType::num_bytes() const {
+  auto num_bytes_per_element =
+      CallTemplatedFunctorWithResult<GetNumBytes, size_t>(element_type());
+  return num_bytes_per_element * shape().num_elements();
+}
+
+bool Tensor::operator==(const Tensor& other) const {
+  return (type() == other.type()) &&
+         !std::memcmp(buffer(), other.buffer(), num_bytes());
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+size_t QuantizedTensorType::num_bytes() const {
+  auto num_bytes_per_element =
+      CallTemplatedFunctorWithResult<GetNumBytes, size_t>(
+          element_type().storage_type());
+  return num_bytes_per_element * shape_.num_elements();
+}
+
+bool QuantizedTensor::operator==(const QuantizedTensor& other) const {
+  return (type() == other.type()) &&
+         !std::memcmp(buffer(), other.buffer(), num_bytes());
+}
+
+QuantizedTensorType QuantizedTensor::baseline_type() const {
+  // For quantized types, we ignore scales and zero points.
+  auto shape = type_.shape();
+
+  if (is_per_tensor_quantized()) {
+    QuantizedTensorElementType element_type(
+        storage_type(), expressed_type(),
+        QuantizedParameter{.scale = 1.0, .zero_point = 0},
+        type_.element_type().storage_min(), type_.element_type().storage_max());
+    return QuantizedTensorType(std::move(shape), std::move(element_type));
+
+  } else {
+    auto n = shape.dim(*type_.element_type().quantized_dimension());
+    std::vector<QuantizedParameter> parameters(n,
+                                               {.scale = 1.0, .zero_point = 0});
+    QuantizedTensorElementType element_type(
+        storage_type(), expressed_type(), std::move(parameters),
+        type_.element_type().storage_min(), type_.element_type().storage_max());
+    return QuantizedTensorType(std::move(shape), std::move(element_type));
+  }
+}
+
+}  // namespace stablehlo

tensorflow/lite/experimental/shlo/legacy/src/storage.h

@@ -0,0 +1,124 @@
+/* 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_SRC_STORAGE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_STORAGE_H_
+
+#include <cstddef>
+#include <cstdint>
+
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/bf16.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/f16.h"
+
+namespace stablehlo {
+
+template <ElementType element_type>
+struct Storage;
+
+template <>
+struct Storage<ElementType::kI1> {
+  using Type = uint8_t;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kSI8> {
+  using Type = int8_t;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kSI16> {
+  using Type = int16_t;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kSI32> {
+  using Type = int32_t;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kBF16> {
+  using Type = BF16;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kF16> {
+  using Type = F16;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+template <>
+struct Storage<ElementType::kF32> {
+  using Type = float;
+  static Type Get(const void* buffer, size_t idx) {
+    auto p = static_cast<const Type*>(buffer);
+    return p[idx];
+  }
+  static void Set(void* buffer, size_t idx, Type value) {
+    auto p = static_cast<Type*>(buffer);
+    p[idx] = value;
+  }
+};
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_STORAGE_H_

tensorflow/lite/experimental/shlo/legacy/src/uniform_dequantize_quantize.cc

@@ -0,0 +1,130 @@
+/* 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 "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/util.h"
+
+namespace stablehlo {
+
+namespace {
+
+absl::Status CheckDequantizeParameters(const QuantizedTensor& operand,
+                                       Tensor& result) {
+  if (operand.shape() != result.shape()) {
+    return absl::InvalidArgumentError("Inconsistent input/output shapes");
+  } else if (operand.expressed_type() != result.element_type()) {
+    return absl::InvalidArgumentError("Inconsistent element types");
+  } else if (!operand.is_per_tensor_quantized()) {
+    return absl::InvalidArgumentError("Unsupported input quantization");
+  }
+
+  if (operand.layout().has_strides() || result.layout().has_strides()) {
+    return absl::InvalidArgumentError("Stides not supported yet");
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+absl::Status UniformDequantize(const QuantizedTensor& operand, Tensor& result) {
+  if (auto check = CheckDequantizeParameters(operand, result); !check.ok()) {
+    return check;
+  }
+
+  const QuantizedParameter& operand_quant_param =
+      operand.type().element_type().parameters(0);
+
+  size_t n = operand.num_elements();
+
+  using S = Storage<storage_type>;
+  using E = Storage<expressed_type>;
+  auto operand_buffer = operand.buffer();
+  auto result_buffer = result.buffer();
+
+  for (size_t i = 0; i < n; ++i) {
+    auto operand_storage = S::Get(operand_buffer, i);
+    auto operand_expressed = Dequantize<storage_type, expressed_type>(
+        operand_storage, operand_quant_param);
+    auto result_expressed = operand_expressed;
+    E::Set(result_buffer, i, result_expressed);
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status CheckQuantizeParameters(const Tensor& operand,
+                                     QuantizedTensor& result) {
+  if (operand.shape() != result.shape()) {
+    return absl::InvalidArgumentError("Inconsistent input/output shapes");
+  } else if (operand.element_type() != result.expressed_type()) {
+    return absl::InvalidArgumentError("Inconsistent element types");
+  } else if (!result.is_per_tensor_quantized()) {
+    return absl::InvalidArgumentError("Unsupported output quantization");
+  }
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+absl::Status UniformQuantize(const Tensor& operand, QuantizedTensor& result) {
+  if (auto check = CheckQuantizeParameters(operand, result); !check.ok()) {
+    return check;
+  }
+
+  const QuantizedParameter& result_quant_param =
+      result.type().element_type().parameters(0);
+
+  size_t n = operand.num_elements();
+
+  using S = Storage<storage_type>;
+  using E = Storage<expressed_type>;
+  auto operand_buffer = operand.buffer();
+  auto result_buffer = result.buffer();
+
+  using ET = typename E::Type;
+  ET result_scale_inv = ET(1.0) / static_cast<ET>(result_quant_param.scale);
+
+  for (size_t i = 0; i < n; ++i) {
+    auto operand_expressed = E::Get(operand_buffer, i);
+    auto result_expressed = operand_expressed;
+    auto result_storage = QuantizePartial<storage_type, expressed_type>(
+        result_expressed, result_scale_inv, result_quant_param.zero_point);
+    S::Set(result_buffer, i, result_storage);
+  }
+
+  if (auto status = CompleteQuantization<storage_type>(result); !status.ok()) {
+    return status;
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status UniformDequantize(const QuantizedTensor& operand, Tensor& result) {
+  DISPATCH_QUANTIZED(UniformDequantize, operand.storage_type(),
+                     operand.expressed_type(), operand, result);
+}
+
+absl::Status UniformQuantize(const Tensor& operand, QuantizedTensor& result) {
+  DISPATCH_QUANTIZED(UniformQuantize, result.storage_type(),
+                     result.expressed_type(), operand, result);
+}
+
+}  // namespace stablehlo

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

@@ -0,0 +1,221 @@
+/* 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_SRC_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_UTIL_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <optional>
+#include <ostream>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/legacy/include/shlo.h"
+#include "tensorflow/lite/experimental/shlo/legacy/src/storage.h"
+
+namespace stablehlo {
+
+template <ElementType storage_type, ElementType expressed_type>
+inline typename Storage<expressed_type>::Type Dequantize(
+    typename Storage<storage_type>::Type quantized_value,
+    const QuantizedParameter& quant_param) {
+  using ST = typename Storage<storage_type>::Type;
+  using ET = typename Storage<expressed_type>::Type;
+  auto sub = (quantized_value - static_cast<ST>(quant_param.zero_point));
+  return static_cast<ET>(sub) * static_cast<ET>(quant_param.scale);
+}
+
+template <ElementType storage_type, ElementType expressed_type>
+inline typename Storage<storage_type>::Type QuantizePartial(
+    typename Storage<expressed_type>::Type expressed_value,
+    typename Storage<expressed_type>::Type scale_inv,
+    typename Storage<storage_type>::Type zero_point) {
+  using ST = typename Storage<storage_type>::Type;
+  using ET = typename Storage<expressed_type>::Type;
+  ET rounding_extra = (expressed_value >= 0) ? ET(0.5) : ET(-0.5);
+  ET tmp = (expressed_value * scale_inv + rounding_extra);
+  // Clamp the value in case of overflow/underflow. This is needed to avoid
+  // getting a SIGILL exception when casting down below.
+  ET max = std::numeric_limits<ST>::max();
+  ET min = std::numeric_limits<ST>::min();
+  if (tmp > max) {
+    tmp = max;
+  } else if (tmp < min) {
+    tmp = min;
+  }
+  ST rounded_value = static_cast<ST>(tmp);
+  ST storage_value = rounded_value + zero_point;
+  return storage_value;
+}
+
+template <ElementType storage_type>
+absl::Status CompleteQuantization(void* buffer, size_t n,
+                                  const std::optional<int32_t>& storage_min,
+                                  const std::optional<int32_t>& storage_max) {
+  using S = Storage<storage_type>;
+
+  if (storage_min) {
+    typename S::Type min = *storage_min;
+    for (size_t i = 0; i < n; ++i) {
+      auto storage = S::Get(buffer, i);
+      storage = std::max(storage, min);
+      S::Set(buffer, i, storage);
+    }
+  }
+
+  if (storage_max) {
+    typename S::Type max = *storage_max;
+    for (size_t i = 0; i < n; ++i) {
+      auto storage = S::Get(buffer, i);
+      storage = std::min(storage, max);
+      S::Set(buffer, i, storage);
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+template <ElementType storage_type>
+absl::Status CompleteQuantization(QuantizedTensor& result) {
+  if (storage_type != result.storage_type()) {
+    return absl::InvalidArgumentError("Unexpected storage type");
+  }
+
+  size_t n = result.num_elements();
+  auto result_buffer = result.buffer();
+  const auto& result_storage_min = result.type().element_type().storage_min();
+  const auto& result_storage_max = result.type().element_type().storage_max();
+  return CompleteQuantization<storage_type>(
+      result_buffer, n, result_storage_min, result_storage_max);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+template <ElementType storage_type, ElementType expressed_type, typename Op>
+inline typename Storage<storage_type>::Type DequantizeOpQuantizePartial(
+    typename Storage<expressed_type>::Type operand_storage,
+    const QuantizedParameter& operand_quant_param,
+    typename Storage<expressed_type>::Type result_scale_inv,
+    typename Storage<storage_type>::Type result_zero_point, Op&& op) {
+  auto operand_expressed = Dequantize<storage_type, expressed_type>(
+      operand_storage, operand_quant_param);
+  auto result_expressed = op(operand_expressed);
+  return QuantizePartial<storage_type, expressed_type>(
+      result_expressed, result_scale_inv, result_zero_point);
+}
+
+template <ElementType storage_type, ElementType expressed_type, typename Op>
+inline typename Storage<storage_type>::Type DequantizeOpQuantizePartial(
+    typename Storage<expressed_type>::Type lhs_storage,
+    typename Storage<expressed_type>::Type rhs_storage,
+    const QuantizedParameter& lhs_quant_param,
+    const QuantizedParameter& rhs_quant_param,
+    typename Storage<expressed_type>::Type result_scale_inv,
+    typename Storage<storage_type>::Type result_zero_point, Op&& op) {
+  auto lhs_expressed =
+      Dequantize<storage_type, expressed_type>(lhs_storage, lhs_quant_param);
+  auto rhs_expressed =
+      Dequantize<storage_type, expressed_type>(rhs_storage, rhs_quant_param);
+  auto result_expressed = op(lhs_expressed, rhs_expressed);
+  return QuantizePartial<storage_type, expressed_type>(
+      result_expressed, result_scale_inv, result_zero_point);
+}
+
+// /////////////////////////////////////////////////////////////////////////////
+
+class TensorIndex {
+ public:
+  explicit TensorIndex(const Shape& shape)
+      : shape_(shape), index_(shape.rank(), 0), linear_index_(0) {}
+  auto operator[](size_t idx) const { return index_[idx]; }
+  void set(size_t idx, DimensionSize value) {
+    index_[idx] = value;
+    linear_index_.reset();
+  }
+  void set(const TensorIndex& other) {
+    index_ = other.index_;
+    linear_index_.reset();
+  }
+  // Return a linearized index assuming the shape's dimension 0 is the major
+  // index (i.e., slowest moving dimension) and the shape's dimension R-1, where
+  // R is the rank, is the minor index (i.e., the fastest moving dimension). For
+  // instance, for tensor<2x3xf32> dimension 1 (of size 3) is the fastest moving
+  // dimension.
+  size_t linearize() const {
+    if (!linear_index_) {
+      linear_index_ = compute_linear_index();
+    }
+    return *linear_index_;
+  }
+
+ private:
+  friend class TensorIndexIterator;
+  friend std::ostream& operator<<(std::ostream&, const TensorIndex&);
+
+  size_t compute_linear_index() const {
+    auto n = index_.size();
+    size_t linear_index = 0;
+    for (auto i = 0; i < n; ++i) {
+      linear_index = (linear_index * shape_.dim(i)) + index_[i];
+    }
+    return linear_index;
+  }
+
+  bool advance() {
+    auto n = index_.size();
+    index_[n - 1]++;
+    if (linear_index_) {
+      (*linear_index_)++;
+    }
+    for (int i = n - 1; i >= 0; --i) {
+      if (index_[i] == shape_.dim(i)) {
+        if ((i - 1) >= 0) {
+          index_[i] = 0;
+          index_[i - 1]++;
+        } else {
+          // Overflow.
+          return false;
+        }
+      }
+    }
+    return true;
+  }
+
+  Shape shape_;
+  std::vector<DimensionSize> index_;
+  mutable std::optional<size_t> linear_index_;
+};
+
+class TensorIndexIterator {
+ public:
+  explicit TensorIndexIterator(const Shape& shape) : index_(shape) {}
+  TensorIndexIterator& operator++() {
+    has_next_ = index_.advance();
+    return *this;
+  }
+  const TensorIndex& operator*() const { return index_; }
+  const TensorIndex* operator->() const { return &index_; }
+  bool has_next() const { return has_next_; }
+
+ private:
+  TensorIndex index_;
+  bool has_next_ = true;
+};
+
+}  // namespace stablehlo
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_LEGACY_SRC_UTIL_H_

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_