chore: import upstream snapshot with attribution

tensorflow/lite/experimental/shlo/BUILD

@@ -0,0 +1,269 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+
+# StableHLO Reference Library
+load("build_def.bzl", "shlo_ref_linkopts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "shlo",
+    deps = [
+        ":tensor",
+    ],
+)
+
+cc_library(
+    name = "tensor",
+    srcs = ["tensor.cc"],
+    hdrs = ["tensor.h"],
+    deps = [
+        ":data_type",
+        ":overload",
+        ":quantized_tensor_element_type",
+        ":shape",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "i4",
+    hdrs = ["i4.h"],
+    deps = [],
+)
+
+cc_test(
+    name = "i4_test",
+    srcs = ["i4_test.cc"],
+    deps = [
+        ":i4",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "tensor_test",
+    srcs = ["tensor_test.cc"],
+    deps = [
+        ":data_type",
+        ":tensor",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "shape",
+    srcs = ["shape.cc"],
+    hdrs = ["shape.h"],
+    deps = [
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "shape_test",
+    srcs = ["shape_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":shape",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "quantized_tensor_element_type",
+    srcs = ["quantized_tensor_element_type.cc"],
+    hdrs = ["quantized_tensor_element_type.h"],
+    deps = [
+        ":data_type",
+        ":shape",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_test(
+    name = "quantized_tensor_element_type_test",
+    srcs = ["quantized_tensor_element_type_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":data_type",
+        ":quantized_tensor_element_type",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "bf16",
+    hdrs = ["bf16.h"],
+    deps = [
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/log:absl_check",
+    ],
+)
+
+cc_test(
+    name = "bf16_test",
+    srcs = ["bf16_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":bf16",
+        "@com_google_absl//absl/base",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "f16",
+    hdrs = ["f16.h"],
+    deps = ["@FP16"],
+)
+
+cc_library(
+    name = "f16_emulated",
+    hdrs = ["f16.h"],
+    copts = ["-DSHLO_REF_EMULATE_F16=1"],
+    deps = ["@FP16"],
+)
+
+cc_test(
+    name = "f16_test",
+    srcs = ["f16_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":f16",
+        "@com_google_absl//absl/base",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "f16_emulated_test",
+    srcs = ["f16_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":f16_emulated",
+        "@com_google_absl//absl/base",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "data_type",
+    hdrs = ["data_type.h"],
+    deps = [
+        ":bf16",
+        ":f16",
+        ":i4",
+    ],
+)
+
+cc_library(
+    name = "dispatch",
+    hdrs = ["dispatch.h"],
+    visibility = ["//tensorflow/lite/experimental/shlo:__subpackages__"],
+)
+
+cc_test(
+    name = "dispatch_test",
+    srcs = ["dispatch_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":dispatch",
+        ":status_matcher",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "overload",
+    hdrs = ["overload.h"],
+    visibility = ["//tensorflow/lite/experimental/shlo:__subpackages__"],
+)
+
+cc_test(
+    name = "overload_test",
+    srcs = ["overload_test.cc"],
+    deps = [
+        ":overload",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "quantize",
+    hdrs = ["quantize.h"],
+    deps = [":data_type"],
+)
+
+cc_test(
+    name = "quantize_test",
+    srcs = ["quantize_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":data_type",
+        ":quantize",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "status_matcher",
+    testonly = True,
+    hdrs = ["status_matcher.h"],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "tensor_with_data",
+    testonly = True,
+    hdrs = ["tensor_with_data.h"],
+    deps = [
+        ":data_type",
+        ":quantize",
+        ":quantized_tensor_element_type",
+        ":shape",
+        ":tensor",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+cc_library(
+    name = "tensor_matcher",
+    testonly = True,
+    hdrs = ["tensor_matcher.h"],
+    deps = [
+        ":data_type",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_test(
+    name = "tensor_matcher_test",
+    srcs = ["tensor_matcher_test.cc"],
+    linkopts = shlo_ref_linkopts(),
+    deps = [
+        ":data_type",
+        ":shape",
+        ":tensor_matcher",
+        ":tensor_with_data",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/lite/experimental/shlo/README.md

@@ -0,0 +1,230 @@
+# StableHLO C++ Reference Library
+
+The goal of this library is to provide a C++ reference implementation of
+StableHLO kernels.
+
+## Contributing
+
+Please review the [Tensorflow Contributing Guide] for the repository's
+contributing guidelines.
+
+The code makes use of C++17 and is built using Bazel.
+
+Unless specified, the [Google style guide] should be followed. Clang-format with
+`google` style should be used for automatic code formatting.
+
+To keep familiarity for people who are used to working with StableHLO, the data
+structures try to follow the naming and hierarchy that are found in the
+[StableHLO specification][stablehlo]
+
+While the library does not strive for performance, we try to avoid unnecessary
+performance penalties. This means avoiding dynamic allocation when possible of
+moving use cases to the `Create` or `Prepare` functions (in order of
+preference).
+
+### Adding an Operation
+
+Refer to the [specification][stablehlo-op] for the naming of an operation, its
+attributes and its inputs.
+
+#### API
+
+An operation is defined using a state structure and three functions.
+
+-   `ExampleOp` is the class/structure that keeps the operation state. It
+    defines a public (possibly empty) `Attributes` structure that holds the
+    attributes described in the operation specification.
+
+    > Tip: Search for `Input attributes` in the [specification][stablehlo] for
+    > more information about attributes.
+
+    > Tip: When reading the specification, the difference between input
+    > attributes and input values is not immediately apparent. Check out the
+    > examples that are given to distinguish them. The definitive authority is
+    > the [StableHLO dialect definition][stablehlo-dialect]: check out the
+    > operations `arguments` declaration for `*Attr` input types.
+
+    ```cpp
+    // Operation data.
+    class ExampleOp {
+     public:
+      // The attributes are a direct mapping of the StableHLO spec.
+      struct Attributes {
+        int64_t attribute_one;
+        float attribute_two;
+      };
+    };
+    ```
+
+-   `Create` initialises the operation data using its attributes as passed
+    through the Attributes structure.
+
+    ```cpp
+    ExampleOp Create(const ExampleOp::Attributes&);
+    ```
+
+-   `Prepare` sets up data and pre-computations that should be reused between
+    evaluations. **In case of dynamic tensors, this step also computes the
+    output tensor dimensions** and should set them.
+
+    -   Preconditions:
+        -   Input tensor shapes are known.
+    -   Postconditions:
+        -   Output tensor shapes are set and valid.
+
+    ```cpp
+    // When an unknown number of tensors can be passed.
+    Status Prepare(ExampleOp& op, const absl::Span<Tensor>& inputs, absl::Span<Tensor>& outputs);
+    // When the number of input/output tensors is known at compile time we can provide an overload
+    Status Prepare(ExampleOp& op, const Tensor& lhs, const Tensor& rhs, Tensor& output);
+    ```
+
+-   `Evaluate` computes the operation result.
+
+    -   Preconditions:
+        -   Input tensor shapes are the same as what was passed to Prepare.
+        -   Input tensor data is known.
+        -   Output tensor shape is known.
+        -   Output tensor buffer is set and allocated.
+    -   Postconditions:
+        -   Output tensor buffers are filled with the operation result.
+
+    ```cpp
+    // When an unknown number of tensors can be passed.
+    Status Eval(ExampleOp& op, const absl::Span<Tensor>& inputs, absl::Span<Tensor>& outputs);
+    // When the number of input/output tensors is known at compile time.
+    Status Eval(ExampleOp& op, const Tensor& lhs, const Tensor& rhs, Tensor& output);
+    ```
+
+Specific operations may define extra functions for implementation configuration
+or tweaks.
+
+#### Bazel
+
+Each operation should be defined in a separate library with the associated tests
+and benchmarks. The code should live in the `ops` folder.
+
+-   The library name should be the name of the operation in `snake_case`.
+-   The implementation and header files should be the name of the library with
+    the `h/cc` extension.
+
+```bzl
+cc_library(
+  name = "op_name",
+  srcs = [ "op_name.cc" ],
+  hdrs = [ "op_name.h" ],
+  deps = [
+    # ...
+  ]
+)
+```
+
+#### Testing
+
+Testing is done with [GoogleTest]. Each operation should be fully tested for
+result correctness and robustness.
+
+-   The test name should be the name of the library with the `_test` suffix.
+-   Use the result matchers to check for results.
+
+```bzl
+cc_test(
+  name = "op_name_test",
+  srcs = [ "op_name_test.cc" ],
+  hdrs = [ "op_name_test.h" ], # Generally not needed.
+  deps = [
+    # ...
+  ]
+)
+```
+
+#### Benchmarking
+
+Testing is done with [Google Benchmark]. Each operation should be fully tested
+for result correctness and robustness.
+
+-   The benchmark name should be the name of the library with the `_bench`
+    suffix.
+
+```bzl
+cc_test(
+  name = "op_name_bench",
+  srcs = [ "op_name_bench.cc" ],
+  hdrs = [ "op_name_bench.h" ], # Generally not needed.
+  deps = [
+    # ...
+  ]
+)
+```
+
+### Running Tests and Benchmarks
+
+This section is a short introduction to running a binary on device.
+
+#### Useful Flags
+
+The following bazel flags may be useful when benchmarking and debugging.
+
+-   `-c dbg`: Compile in debug mode.
+-   `-c opt`: Compile in optimized mode.
+-   `-gmlt`: Adds line and function name debug information to optimised builds.
+
+#### x86
+
+##### Tests
+
+```sh
+bazel test -c opt --dynamic_mode=off ops:op_name_test
+```
+
+> Note: it is often useful to run test in optimized **and** in debug mode.
+
+##### Benchmarks
+
+```sh
+bazel run -c opt --dynamic_mode=off ops:op_name_bench
+```
+
+#### Android
+
+```sh
+bazel build -c opt --dynamic_mode=off --config=android_arm64 --copt=-DGOOGLE_COMMANDLINEFLAGS_FULL_API=1 ops:op_name_test
+```
+
+Bazel should print the location of the build binary. It should resemble
+`shlo/ops/op_name_test`.
+
+You can then push the binary to the device `/data/local/tmp` folder and run it
+using ADB.
+
+```sh
+adb push shlo/ops/op_name_test /data/local/tmp
+adb shell /data/local/tmp/op_name_test
+```
+
+#### iOS
+
+##### Prerequisites
+
+Follow the instructions for setting up the iOS development environment in the
+TensorFlow Lite [Build for iOS] guide. The `configure` script must be run and
+you must opt-in to iOS development.
+
+##### Building
+
+```
+bazel build -c opt --config=ios_arm64 ops:op_name_test
+```
+
+##### Testing
+
+TODO:
+
+[stablehlo]: https://github.com/openxla/stablehlo/blob/main/docs/spec.md
+[stablehlo-op]: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#operations
+[stablehlo-dialect]: https://github.com/openxla/stablehlo/blob/main/stablehlo/dialect/StablehloOps.td
+[GoogleTest]: https://github.com/google/googletest
+[Google Benchmark]: https://github.com/google/benchmark
+[Google style guide]: https://google.github.io/styleguide/cppguide.html
+[Tensorflow Contributing Guide]: https://github.com/tensorflow/tensorflow/blob/master/CONTRIBUTING.md
+[Build for iOS]: https://www.tensorflow.org/lite/guide/build_ios

tensorflow/lite/experimental/shlo/bf16.h

@@ -0,0 +1,279 @@
+/* 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_BF16_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_BF16_H_
+
+#if defined(__STDCPP_BFLOAT16_T__)
+#include <stdfloat>
+namespace shlo_ref {
+using BF16 = ::std::bfloat16_t;
+}  // namespace shlo_ref
+
+#else
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+#include "absl/base/casts.h"
+#include "absl/log/absl_check.h"
+
+// 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.
+namespace shlo_ref {
+class BF16;
+
+namespace internal {
+BF16 NumericF32ToBF16RoundNearestEven(float v);
+BF16 F32ToBF16RoundNearestEven(float v);
+float BF16ToFloat(BF16 v);
+}  // namespace internal
+}  // namespace shlo_ref
+
+namespace std {
+template <>
+struct numeric_limits< ::shlo_ref::BF16>;
+}  // namespace std
+
+namespace shlo_ref {
+
+class BF16 {
+ public:
+  constexpr BF16() : value_(0) {}
+
+  template <typename T,
+            typename = std::enable_if_t<std::is_convertible_v<T, float> > >
+  explicit BF16(T x) {
+    if constexpr (std::is_same_v<T, bool>) {
+      value_ = static_cast<uint16_t>(x) * 0x3f80;
+    } else if constexpr (std::numeric_limits<T>::is_integer) {
+      *this = internal::NumericF32ToBF16RoundNearestEven(static_cast<float>(x));
+    } else {
+      *this = internal::F32ToBF16RoundNearestEven(static_cast<float>(x));
+    }
+  }
+
+  // Tagged constructor to allow construction from bits
+  struct bitcast_construct_t {};
+  explicit constexpr BF16(bitcast_construct_t, uint16_t value)
+      : value_(value) {}
+
+  // Lossless conversion to `float`.
+  operator float() const {  // NOLINT: Allow implicit conversions to float.
+    return internal::BF16ToFloat(*this);
+  }
+
+  // Assignment operators
+  BF16& operator=(float v) { return *this = static_cast<BF16>(v); }
+  BF16& operator=(bool v) { return *this = static_cast<BF16>(v); }
+
+  template <typename T>
+  // NOLINTNEXTLINE(misc-unconventional-assign-operator)
+  std::enable_if_t<std::numeric_limits<T>::is_integer, BF16&> operator=(T v) {
+    return *this = static_cast<BF16>(v);
+  }
+
+#define INTERNAL_BF16_ARITHMETIC_OP(OP)                          \
+  friend BF16 operator OP(BF16 x, BF16 y) {                      \
+    return BF16(static_cast<float>(x) OP static_cast<float>(y)); \
+  }
+
+#define INTERNAL_BF16_ARITHMETIC_ASSIGN_OP(OP)                       \
+  friend BF16& operator OP##=(BF16 & x, BF16 y) {                    \
+    return x = BF16(static_cast<float>(x) OP static_cast<float>(y)); \
+  }
+
+  INTERNAL_BF16_ARITHMETIC_OP(+)
+  INTERNAL_BF16_ARITHMETIC_ASSIGN_OP(+)
+  INTERNAL_BF16_ARITHMETIC_OP(-)
+  INTERNAL_BF16_ARITHMETIC_ASSIGN_OP(-)
+  INTERNAL_BF16_ARITHMETIC_OP(*)
+  INTERNAL_BF16_ARITHMETIC_ASSIGN_OP(*)
+  INTERNAL_BF16_ARITHMETIC_OP(/)
+  INTERNAL_BF16_ARITHMETIC_ASSIGN_OP(/)
+  INTERNAL_BF16_ARITHMETIC_OP(==)
+  INTERNAL_BF16_ARITHMETIC_OP(!=)
+  INTERNAL_BF16_ARITHMETIC_OP(<)
+  INTERNAL_BF16_ARITHMETIC_OP(<=)
+  INTERNAL_BF16_ARITHMETIC_OP(>)
+  INTERNAL_BF16_ARITHMETIC_OP(>=)
+
+#undef INTERNAL_BF16_ARITHMETIC_OP
+#undef INTERNAL_BF16_ARITHMETIC_ASSIGN_OP
+
+  // Unary negation.
+  friend BF16 operator-(BF16 x) {
+    BF16 result;
+    result.value_ = x.value_ ^ 0x8000;
+    return result;
+  }
+
+  // Unary plus
+  friend BF16 operator+(BF16 x) { return x; }
+
+ private:
+  uint16_t value_;
+};
+
+inline bool isinf(BF16 x) { return std::isinf(static_cast<float>(x)); }
+inline bool signbit(BF16 x) { return std::signbit(static_cast<float>(x)); }
+inline bool isnan(BF16 x) { return std::isnan(static_cast<float>(x)); }
+inline bool isfinite(BF16 x) { return std::isfinite(static_cast<float>(x)); }
+inline BF16 abs(BF16 x) { return BF16(std::abs(static_cast<float>(x))); }
+inline BF16 exp(BF16 x) { return BF16(std::exp(static_cast<float>(x))); }
+inline BF16 exp2(BF16 x) { return BF16(std::exp2(static_cast<float>(x))); }
+inline BF16 expm1(BF16 x) { return BF16(std::expm1(static_cast<float>(x))); }
+inline BF16 log(BF16 x) { return BF16(std::log(static_cast<float>(x))); }
+inline BF16 log1p(BF16 x) { return BF16(std::log1p(static_cast<float>(x))); }
+inline BF16 log10(BF16 x) { return BF16(std::log10(static_cast<float>(x))); }
+inline BF16 log2(BF16 x) { return BF16(std::log2(static_cast<float>(x))); }
+inline BF16 sqrt(BF16 x) { return BF16(std::sqrt(static_cast<float>(x))); }
+inline BF16 pow(BF16 x, BF16 y) {
+  return BF16(std::pow(static_cast<float>(x), static_cast<float>(y)));
+}
+inline BF16 sin(BF16 x) { return BF16(std::sin(static_cast<float>(x))); }
+inline BF16 cos(BF16 x) { return BF16(std::cos(static_cast<float>(x))); }
+inline BF16 tan(BF16 x) { return BF16(std::tan(static_cast<float>(x))); }
+inline BF16 asin(BF16 x) { return BF16(std::asin(static_cast<float>(x))); }
+inline BF16 acos(BF16 x) { return BF16(std::acos(static_cast<float>(x))); }
+inline BF16 atan(BF16 x) { return BF16(std::atan(static_cast<float>(x))); }
+inline BF16 sinh(BF16 x) { return BF16(std::sinh(static_cast<float>(x))); }
+inline BF16 cosh(BF16 x) { return BF16(std::cosh(static_cast<float>(x))); }
+inline BF16 tanh(BF16 x) { return BF16(std::tanh(static_cast<float>(x))); }
+inline BF16 asinh(BF16 x) { return BF16(std::asinh(static_cast<float>(x))); }
+inline BF16 acosh(BF16 x) { return BF16(std::acosh(static_cast<float>(x))); }
+inline BF16 atanh(BF16 x) { return BF16(std::atanh(static_cast<float>(x))); }
+inline BF16 floor(BF16 x) { return BF16(std::floor(static_cast<float>(x))); }
+inline BF16 trunc(BF16 x) { return BF16(std::trunc(static_cast<float>(x))); }
+inline BF16 rint(BF16 x) { return BF16(std::rint(static_cast<float>(x))); }
+inline BF16 ceil(BF16 x) { return BF16(std::ceil(static_cast<float>(x))); }
+inline BF16 fmod(BF16 x, BF16 y) {
+  return BF16(std::fmod(static_cast<float>(x), static_cast<float>(y)));
+}
+inline BF16 fmin(BF16 a, BF16 b) {
+  return BF16(std::fmin(static_cast<float>(a), static_cast<float>(b)));
+}
+inline BF16 fmax(BF16 a, BF16 b) {
+  return BF16(std::fmax(static_cast<float>(a), static_cast<float>(b)));
+}
+
+namespace internal {
+
+inline BF16 NumericF32ToBF16RoundNearestEven(float v) {
+  ABSL_CHECK(!std::isnan(v));
+
+  uint32_t input = absl::bit_cast<uint32_t>(v);
+  const uint32_t lsb = (input >> 16) & 1;
+  const uint32_t rounding_bias = 0x7fff + lsb;
+  input += rounding_bias;
+  return absl::bit_cast<BF16, uint16_t>(input >> 16);
+}
+
+inline BF16 F32ToBF16RoundNearestEven(float v) {
+  if (std::isnan(v)) {
+    return BF16(BF16::bitcast_construct_t{},
+                static_cast<uint16_t>(
+                    (absl::bit_cast<uint32_t>(v) | 0x00200000u) >> 16));
+  }
+  return NumericF32ToBF16RoundNearestEven(v);
+}
+
+inline float BF16ToFloat(BF16 v) {
+  return absl::bit_cast<float>(
+      static_cast<uint32_t>(absl::bit_cast<uint16_t>(v)) << 16);
+}
+}  // namespace internal
+}  // namespace shlo_ref
+
+// Specialized std::numeric_limits for BF16
+namespace std {
+template <>
+class numeric_limits<shlo_ref::BF16> {
+ public:
+  static constexpr bool is_specialized = true;      // NOLINT
+  static constexpr bool is_signed = true;           // NOLINT
+  static constexpr bool is_integer = false;         // NOLINT
+  static constexpr bool is_exact = false;           // NOLINT
+  static constexpr bool has_infinity = true;        // NOLINT
+  static constexpr bool has_quiet_NaN = true;       // NOLINT
+  static constexpr bool has_signaling_NaN = true;   // NOLINT
+  static constexpr float_denorm_style has_denorm =  // NOLINT
+      std::denorm_present;
+  static constexpr bool has_denorm_loss = false;    // NOLINT
+  static constexpr float_round_style round_style =  // NOLINT
+      numeric_limits<float>::round_style;
+  static constexpr bool is_iec559 = true;   // NOLINT
+  static constexpr bool is_bounded = true;  // NOLINT
+  static constexpr bool is_modulo = false;  // NOLINT
+  static constexpr int digits = 8;          // NOLINT
+  static constexpr int digits10 = 2;        // NOLINT
+  static constexpr int max_digits10 = 4;    // NOLINT
+  static constexpr int radix = 2;           // NOLINT
+  static constexpr int min_exponent =       // NOLINT
+      numeric_limits<float>::min_exponent;
+  static constexpr int min_exponent10 =  // NOLINT
+      numeric_limits<float>::min_exponent10;
+  static constexpr int max_exponent =  // NOLINT
+      numeric_limits<float>::max_exponent;
+  static constexpr int max_exponent10 =  // NOLINT
+      numeric_limits<float>::max_exponent10;
+  static constexpr bool traps = numeric_limits<float>::traps;  // NOLINT
+  static constexpr bool tinyness_before =                      // NOLINT
+      numeric_limits<float>::tinyness_before;
+
+  static constexpr shlo_ref::BF16(min)() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x0080));
+  }
+  static constexpr shlo_ref::BF16 lowest() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0xff7f));
+  }
+  static constexpr shlo_ref::BF16(max)() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x7f7f));
+  }
+  static constexpr shlo_ref::BF16 epsilon() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x3c00));
+  }
+  static constexpr shlo_ref::BF16 round_error() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x3f00));
+  }
+  static constexpr shlo_ref::BF16 infinity() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x7f80));
+  }
+  static constexpr shlo_ref::BF16 quiet_NaN() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x7fc0));
+  }
+  static constexpr shlo_ref::BF16 signaling_NaN() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x7f81));
+  }
+  static constexpr shlo_ref::BF16 denorm_min() {
+    return shlo_ref::BF16(shlo_ref::BF16::bitcast_construct_t{},
+                          static_cast<uint16_t>(0x0001));
+  }
+};
+}  // namespace std
+
+#endif
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_BF16_H_

tensorflow/lite/experimental/shlo/bf16_test.cc

@@ -0,0 +1,340 @@
+/* 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/bf16.h"
+
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+
+namespace shlo_ref {
+namespace {
+
+::testing::Matcher<BF16> MatchesBits(uint16_t bits) {
+  return ::testing::ResultOf([](BF16 y) { return absl::bit_cast<uint16_t>(y); },
+                             ::testing::Eq(bits));
+}
+
+::testing::Matcher<float> NearFloat(float x, float relative_error = 1e-3) {
+  return ::testing::FloatNear(x, std::abs(x) * relative_error);
+}
+
+float BinaryToFloat(uint32_t sign, uint32_t exponent, uint32_t high_mantissa,
+                    uint32_t low_mantissa) {
+  float dest;
+  uint32_t src =
+      (sign << 31) + (exponent << 23) + (high_mantissa << 16) + low_mantissa;
+  memcpy(static_cast<void*>(&dest), static_cast<const void*>(&src),
+         sizeof(dest));
+  return dest;
+}
+
+template <typename T>
+void TestRoundtrips() {
+  for (T value : {
+           -std::numeric_limits<T>::infinity(),
+           std::numeric_limits<T>::infinity(),
+           T(-1.0),
+           T(-0.5),
+           T(-0.0),
+           T(1.0),
+           T(0.5),
+           T(0.0),
+       }) {
+    EXPECT_EQ(value, static_cast<T>(static_cast<BF16>(value)));
+  }
+}
+
+TEST(BF16Test, FloatRoundtrips) { TestRoundtrips<float>(); }
+
+TEST(BF16Test, DoubleRoundtrips) { TestRoundtrips<double>(); }
+
+TEST(BF16Test, Float16Roundtrips) { TestRoundtrips<BF16>(); }
+
+TEST(BF16Test, ConversionFromFloat) {
+  EXPECT_THAT(BF16(1.0f), MatchesBits(0x3f80));
+  EXPECT_THAT(BF16(0.5f), MatchesBits(0x3f00));
+  EXPECT_THAT(BF16(0.33333f), MatchesBits(0x3eab));
+  EXPECT_THAT(BF16(3.38e38f), MatchesBits(0x7f7e));
+  EXPECT_THAT(BF16(3.40e38f), MatchesBits(0x7f80));  // Becomes infinity.
+}
+
+TEST(BF16Test, RoundToNearestEven) {
+  float val1 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c00}));
+  float val2 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c01}));
+  float val3 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c02}));
+  EXPECT_THAT(BF16(0.5f * (val1 + val2)), MatchesBits(0x3c00));
+  EXPECT_THAT(BF16(0.5f * (val2 + val3)), MatchesBits(0x3c02));
+}
+
+TEST(BF16Test, ConversionFromInt) {
+  EXPECT_THAT(BF16(-1), MatchesBits(0xbf80));
+  EXPECT_THAT(BF16(0), MatchesBits(0x0000));
+  EXPECT_THAT(BF16(1), MatchesBits(0x3f80));
+  EXPECT_THAT(BF16(2), MatchesBits(0x4000));
+  EXPECT_THAT(BF16(3), MatchesBits(0x4040));
+  EXPECT_THAT(BF16(12), MatchesBits(0x4140));
+}
+
+TEST(BF16Test, ConversionFromBool) {
+  EXPECT_THAT(BF16(false), MatchesBits(0x0000));
+  EXPECT_THAT(BF16(true), MatchesBits(0x3f80));
+}
+
+TEST(BF16Test, ConversionToBool) {
+  EXPECT_EQ(static_cast<bool>(BF16(3)), true);
+  EXPECT_EQ(static_cast<bool>(BF16(0.33333f)), true);
+  EXPECT_EQ(BF16(-0.0), false);
+  EXPECT_EQ(static_cast<bool>(BF16(0.0)), false);
+}
+
+TEST(BF16Test, ExplicitConversionToFloat) {
+  EXPECT_EQ(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x0000)), 0.0f);
+  EXPECT_EQ(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x3f80)), 1.0f);
+}
+
+TEST(BF16Test, ImplicitConversionToFloat) {
+  EXPECT_EQ((absl::bit_cast<BF16, uint16_t>(0x0000)), 0.0f);
+  EXPECT_EQ((absl::bit_cast<BF16, uint16_t>(0x3f80)), 1.0f);
+}
+
+TEST(BF16Test, Zero) {
+  EXPECT_EQ(BF16(0.0f), BF16(0.0f));
+  EXPECT_EQ(BF16(-0.0f), BF16(0.0f));
+  EXPECT_EQ(BF16(-0.0f), BF16(-0.0f));
+  EXPECT_THAT(BF16(0.0f), MatchesBits(0x0000));
+  EXPECT_THAT(BF16(-0.0f), MatchesBits(0x8000));
+}
+
+TEST(BF16Test, DefaultConstruct) {
+  EXPECT_EQ(static_cast<float>(BF16()), 0.0f);
+}
+
+TEST(BF16Test, Conversion) {
+  for (int i = 0; i < 100; ++i) {
+    float a = i + 1.25;
+    BF16 b = static_cast<BF16>(a);
+    float c = static_cast<float>(b);
+    EXPECT_LE(std::abs(c - a), a / 128);
+  }
+}
+
+TEST(BF16Test, Epsilon) {
+  EXPECT_LE(1.0f, static_cast<float>(std::numeric_limits<BF16>::epsilon() +
+                                     BF16(1.0f)));
+  EXPECT_EQ(1.0f, static_cast<float>(std::numeric_limits<BF16>::epsilon() /
+                                         BF16(2.0f) +
+                                     BF16(1.0f)));
+}
+
+TEST(BF16Test, Negate) {
+  EXPECT_EQ(static_cast<float>(-BF16(3.0f)), -3.0f);
+  EXPECT_EQ(static_cast<float>(-BF16(-4.5f)), 4.5f);
+}
+
+TEST(BF16Test, DivisionByZero) {
+  EXPECT_TRUE(std::isnan(static_cast<float>(BF16(0.0 / 0.0))));
+  EXPECT_TRUE(std::isinf(static_cast<float>(BF16(1.0 / 0.0))));
+  EXPECT_TRUE(std::isinf(static_cast<float>(BF16(-1.0 / 0.0))));
+
+  EXPECT_TRUE(std::isnan(BF16(0.0 / 0.0)));
+  EXPECT_TRUE(std::isinf(BF16(1.0 / 0.0)));
+  EXPECT_TRUE(std::isinf(BF16(-1.0 / 0.0)));
+}
+
+TEST(BF16Test, NonFinite) {
+  EXPECT_FALSE(std::isinf(
+      static_cast<float>(BF16(3.38e38f))));  // Largest finite number.
+  EXPECT_FALSE(std::isnan(static_cast<float>(BF16(0.0f))));
+  EXPECT_TRUE(
+      std::isinf(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0xff80))));
+  EXPECT_TRUE(
+      std::isnan(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0xffc0))));
+  EXPECT_TRUE(
+      std::isinf(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x7f80))));
+  EXPECT_TRUE(
+      std::isnan(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x7fc0))));
+
+  // Exactly same checks as above, just directly on the BF16 representation.
+  EXPECT_FALSE(isinf(absl::bit_cast<BF16, uint16_t>(0x7bff)));
+  EXPECT_FALSE(isnan(absl::bit_cast<BF16, uint16_t>(0x0000)));
+  EXPECT_TRUE(isinf(absl::bit_cast<BF16, uint16_t>(0xff80)));
+  EXPECT_TRUE(isnan(absl::bit_cast<BF16, uint16_t>(0xffc0)));
+  EXPECT_TRUE(isinf(absl::bit_cast<BF16, uint16_t>(0x7f80)));
+  EXPECT_TRUE(isnan(absl::bit_cast<BF16, uint16_t>(0x7fc0)));
+
+  EXPECT_THAT(BF16(BinaryToFloat(0x0, 0xff, 0x40, 0x0)),  // +nan
+              MatchesBits(0x7fe0));
+  EXPECT_THAT(BF16(BinaryToFloat(0x1, 0xff, 0x40, 0x0)),  // -nan
+              MatchesBits(0xffe0));
+}
+
+TEST(BF16Test, NumericLimits) {
+  static_assert(std::numeric_limits<BF16>::is_signed);
+
+  EXPECT_EQ(
+      absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::infinity()),
+      absl::bit_cast<uint16_t>(BF16(std::numeric_limits<float>::infinity())));
+  // There is no guarantee that casting a 32-bit NaN to bfloat16 has a precise
+  // bit pattern.  We test that it is in fact a NaN, then test the signaling
+  // bit (msb of significand is 1 for quiet, 0 for signaling).
+  constexpr uint16_t BFLOAT16_QUIET_BIT = 0x0040;
+  EXPECT_TRUE(isnan(std::numeric_limits<BF16>::quiet_NaN()));
+  EXPECT_TRUE(isnan(BF16(std::numeric_limits<float>::quiet_NaN())));
+  EXPECT_GT((absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::quiet_NaN()) &
+             BFLOAT16_QUIET_BIT),
+            0);
+  EXPECT_GT(
+      (absl::bit_cast<uint16_t>(BF16(std::numeric_limits<float>::quiet_NaN())) &
+       BFLOAT16_QUIET_BIT),
+      0);
+
+  EXPECT_TRUE(isnan(std::numeric_limits<BF16>::signaling_NaN()));
+  EXPECT_TRUE(isnan(BF16(std::numeric_limits<float>::signaling_NaN())));
+  EXPECT_EQ(
+      0, (absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::signaling_NaN()) &
+          BFLOAT16_QUIET_BIT));
+  EXPECT_EQ(0, (absl::bit_cast<uint16_t>(
+                    BF16(std::numeric_limits<float>::signaling_NaN())) &
+                BFLOAT16_QUIET_BIT));
+
+  EXPECT_GT(std::numeric_limits<BF16>::min(), BF16(0.f));
+  EXPECT_GT(std::numeric_limits<BF16>::denorm_min(), BF16(0.f));
+  EXPECT_EQ(std::numeric_limits<BF16>::denorm_min() / BF16(2), BF16(0.f));
+}
+
+TEST(BF16Test, Arithmetic) {
+  EXPECT_EQ(static_cast<float>(BF16(2) + BF16(2)), 4);
+  EXPECT_EQ(static_cast<float>(BF16(2) + BF16(-2)), 0);
+  EXPECT_THAT(static_cast<float>(BF16(0.33333f) + BF16(0.66667f)),
+              NearFloat(1.0f));
+  EXPECT_EQ(static_cast<float>(BF16(2.0f) * BF16(-5.5f)), -11.0f);
+  EXPECT_THAT(static_cast<float>(BF16(1.0f) / BF16(3.0f)), NearFloat(0.3339f));
+  EXPECT_EQ(static_cast<float>(-BF16(4096.0f)), -4096.0f);
+  EXPECT_EQ(static_cast<float>(-BF16(-4096.0f)), 4096.0f);
+}
+
+TEST(BF16Test, Comparison) {
+  EXPECT_TRUE(BF16(1.0f) > BF16(0.5f));
+  EXPECT_TRUE(BF16(0.5f) < BF16(1.0f));
+  EXPECT_FALSE((BF16(1.0f) < BF16(0.5f)));
+  EXPECT_FALSE((BF16(0.5f) > BF16(1.0f)));
+
+  EXPECT_FALSE((BF16(4.0f) > BF16(4.0f)));
+  EXPECT_FALSE((BF16(4.0f) < BF16(4.0f)));
+
+  EXPECT_FALSE((BF16(0.0f) < BF16(-0.0f)));
+  EXPECT_FALSE((BF16(-0.0f) < BF16(0.0f)));
+  EXPECT_FALSE((BF16(0.0f) > BF16(-0.0f)));
+  EXPECT_FALSE((BF16(-0.0f) > BF16(0.0f)));
+
+  EXPECT_TRUE(BF16(0.2f) > BF16(-1.0f));
+  EXPECT_TRUE(BF16(-1.0f) < BF16(0.2f));
+  EXPECT_TRUE(BF16(-16.0f) < BF16(-15.0f));
+
+  EXPECT_TRUE(BF16(1.0f) == BF16(1.0f));
+  EXPECT_TRUE(BF16(1.0f) != BF16(2.0f));
+
+  EXPECT_FALSE((BF16(0.0 / 0.0) == BF16(0.0 / 0.0)));
+  EXPECT_TRUE(BF16(0.0 / 0.0) != BF16(0.0 / 0.0));
+
+  EXPECT_FALSE((BF16(1.0) == BF16(0.0 / 0.0)));
+  EXPECT_FALSE((BF16(1.0) < BF16(0.0 / 0.0)));
+  EXPECT_FALSE((BF16(1.0) > BF16(0.0 / 0.0)));
+  EXPECT_TRUE(BF16(1.0) != BF16(0.0 / 0.0));
+
+  EXPECT_TRUE(BF16(1.0) < BF16(1.0 / 0.0));
+  EXPECT_TRUE(BF16(1.0) > BF16(-1.0 / 0.0));
+}
+
+constexpr float PI = 3.14159265358979323846f;
+
+TEST(BF16Test, BasicFunctions) {
+  // These calls should be found via ADL.
+  EXPECT_EQ(static_cast<float>(abs(BF16(3.5f))), 3.5f);
+  EXPECT_EQ(static_cast<float>(abs(BF16(3.5f))), 3.5f);
+  EXPECT_EQ(static_cast<float>(abs(BF16(-3.5f))), 3.5f);
+  EXPECT_EQ(static_cast<float>(abs(BF16(-3.5f))), 3.5f);
+
+  EXPECT_EQ(static_cast<float>(floor(BF16(3.5f))), 3.0f);
+  EXPECT_EQ(static_cast<float>(floor(BF16(3.5f))), 3.0f);
+  EXPECT_EQ(static_cast<float>(floor(BF16(-3.5f))), -4.0f);
+  EXPECT_EQ(static_cast<float>(floor(BF16(-3.5f))), -4.0f);
+
+  EXPECT_EQ(static_cast<float>(ceil(BF16(3.5f))), 4.0f);
+  EXPECT_EQ(static_cast<float>(ceil(BF16(3.5f))), 4.0f);
+  EXPECT_EQ(static_cast<float>(ceil(BF16(-3.5f))), -3.0f);
+  EXPECT_EQ(static_cast<float>(ceil(BF16(-3.5f))), -3.0f);
+
+  EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(0.0f))), 0.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(0.0f))), 0.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(4.0f))), 2.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(4.0f))), 2.0f);
+
+  EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(0.0f), BF16(1.0f))), 0.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(0.0f), BF16(1.0f))), 0.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(2.0f), BF16(2.0f))), 4.0f);
+  EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(2.0f), BF16(2.0f))), 4.0f);
+
+  EXPECT_EQ(static_cast<float>(exp(BF16(0.0f))), 1.0f);
+  EXPECT_EQ(static_cast<float>(exp(BF16(0.0f))), 1.0f);
+  EXPECT_THAT(static_cast<float>(exp(BF16(PI))),
+              NearFloat(20.f + static_cast<float>(PI)));
+  EXPECT_THAT(static_cast<float>(exp(BF16(PI))),
+              NearFloat(20.f + static_cast<float>(PI)));
+
+  EXPECT_EQ(static_cast<float>(expm1(BF16(0.0f))), 0.0f);
+  EXPECT_EQ(static_cast<float>(expm1(BF16(0.0f))), 0.0f);
+  EXPECT_THAT(static_cast<float>(expm1(BF16(2.0f))), NearFloat(6.375f));
+  EXPECT_THAT(static_cast<float>(expm1(BF16(2.0f))), NearFloat(6.375f));
+
+  EXPECT_EQ(static_cast<float>(log(BF16(1.0f))), 0.0f);
+  EXPECT_EQ(static_cast<float>(log(BF16(1.0f))), 0.0f);
+  EXPECT_THAT(static_cast<float>(log(BF16(10.0f))), NearFloat(2.296875f));
+  EXPECT_THAT(static_cast<float>(log(BF16(10.0f))), NearFloat(2.296875f));
+
+  EXPECT_EQ(static_cast<float>(log1p(BF16(0.0f))), 0.0f);
+  EXPECT_EQ(static_cast<float>(log1p(BF16(0.0f))), 0.0f);
+  EXPECT_THAT(static_cast<float>(log1p(BF16(10.0f))), NearFloat(2.390625f));
+  EXPECT_THAT(static_cast<float>(log1p(BF16(10.0f))), NearFloat(2.390625f));
+}
+
+TEST(BF16Test, TrigonometricFunctions) {
+  EXPECT_THAT(cos(BF16(0.0f)), NearFloat(BF16(std::cos(0.0f))));
+  EXPECT_THAT(cos(BF16(0.0f)), NearFloat(BF16(std::cos(0.0f))));
+  EXPECT_FLOAT_EQ(cos(BF16(PI)), BF16(std::cos(PI)));
+  EXPECT_NEAR(cos(BF16(PI / 2)), BF16(std::cos(PI / 2)), 1e-3);
+  EXPECT_NEAR(cos(BF16(3 * PI / 2)), BF16(std::cos(3 * PI / 2)), 1e-2);
+  EXPECT_THAT(cos(BF16(3.5f)), NearFloat(BF16(std::cos(3.5f))));
+
+  EXPECT_FLOAT_EQ(sin(BF16(0.0f)), BF16(std::sin(0.0f)));
+  EXPECT_FLOAT_EQ(sin(BF16(0.0f)), BF16(std::sin(0.0f)));
+  EXPECT_NEAR(sin(BF16(PI)), BF16(std::sin(PI)), 1e-3);
+  EXPECT_THAT(sin(BF16(PI / 2)), NearFloat(BF16(std::sin(PI / 2))));
+  EXPECT_THAT(sin(BF16(3 * PI / 2)), NearFloat(BF16(std::sin(3 * PI / 2))));
+  EXPECT_THAT(sin(BF16(3.5f)), NearFloat(BF16(std::sin(3.5f))));
+
+  EXPECT_FLOAT_EQ(tan(BF16(0.0f)), BF16(std::tan(0.0f)));
+  EXPECT_FLOAT_EQ(tan(BF16(0.0f)), BF16(std::tan(0.0f)));
+  EXPECT_NEAR(tan(BF16(PI)), BF16(std::tan(PI)), 1e-3);
+  EXPECT_THAT(tan(BF16(3.5f)), NearFloat(BF16(std::tan(3.5f))));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/build_def.bzl

@@ -0,0 +1,12 @@
+"""Build macros for SHLO reference library."""
+
+def shlo_ref_logging_linkopts():
+    """Defines linker flags to enable logging"""
+    return select({
+        "//tensorflow:android": ["-llog"],
+        "//conditions:default": [],
+    })
+
+def shlo_ref_linkopts():
+    """Defines linker flags for linking SHLO binary"""
+    return shlo_ref_logging_linkopts()

tensorflow/lite/experimental/shlo/data_type.h

@@ -0,0 +1,166 @@
+/* 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_DATA_TYPE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_DATA_TYPE_H_
+
+#include <cstdint>
+#include <limits>
+
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/i4.h"
+
+namespace shlo_ref {
+
+// For more information on StableHLO types, see the spec., search for "Element
+// types". The SHLO Device Profile does not include unsigned or 64 bit types.
+enum class DataType {
+  kI1,
+  kSI4,
+  kSI8,
+  kSI16,
+  kSI32,
+  kSI64,
+  kBF16,
+  kF16,
+  kF32,
+};
+
+template <class T>
+struct DefaultStorageDescription {
+  using Type = T;
+  static constexpr Type kMinValue = std::numeric_limits<Type>::lowest();
+  static constexpr Type kMaxValue = std::numeric_limits<Type>::max();
+};
+
+// Storage provides the corresponding C++ type for the given DataType.
+template <DataType data_type>
+struct Storage {};
+
+template <>
+struct Storage<DataType::kI1> : DefaultStorageDescription<bool> {};
+
+template <>
+struct Storage<DataType::kSI4> : DefaultStorageDescription<I4> {};
+
+template <>
+struct Storage<DataType::kSI8> : DefaultStorageDescription<int8_t> {};
+
+template <>
+struct Storage<DataType::kSI16> : DefaultStorageDescription<int16_t> {};
+
+template <>
+struct Storage<DataType::kSI32> : DefaultStorageDescription<int32_t> {};
+
+template <>
+struct Storage<DataType::kSI64> : DefaultStorageDescription<int64_t> {};
+
+template <>
+struct Storage<DataType::kBF16> : DefaultStorageDescription<BF16> {};
+
+template <>
+struct Storage<DataType::kF16> : DefaultStorageDescription<F16> {};
+
+template <>
+struct Storage<DataType::kF32> : DefaultStorageDescription<float> {};
+
+template <DataType data_type>
+using StorageType = typename Storage<data_type>::Type;
+
+constexpr bool IsBool(DataType data_type) { return data_type == DataType::kI1; }
+
+constexpr bool IsSignedInteger(DataType data_type) {
+  return data_type == DataType::kSI4 || data_type == DataType::kSI8 ||
+         data_type == DataType::kSI16 || data_type == DataType::kSI32 ||
+         data_type == DataType::kSI64;
+}
+
+constexpr bool IsUnsignedInteger(DataType data_type) { return false; }
+
+constexpr bool IsInteger(DataType data_type) {
+  return IsSignedInteger(data_type) || IsUnsignedInteger(data_type);
+}
+
+constexpr bool IsFloat(DataType data_type) {
+  return data_type == DataType::kBF16 || data_type == DataType::kF16 ||
+         data_type == DataType::kF32;
+}
+
+template <DataType data_type>
+constexpr int64_t SizeOf() {
+  return sizeof(StorageType<data_type>);
+}
+
+constexpr int64_t SizeOf(DataType data_type) {
+  switch (data_type) {
+    case DataType::kI1:
+      return SizeOf<DataType::kI1>();
+    case DataType::kSI4:
+      return SizeOf<DataType::kSI4>();
+    case DataType::kSI8:
+      return SizeOf<DataType::kSI8>();
+    case DataType::kSI16:
+      return SizeOf<DataType::kSI16>();
+    case DataType::kSI32:
+      return SizeOf<DataType::kSI32>();
+    case DataType::kSI64:
+      return SizeOf<DataType::kSI64>();
+    case DataType::kBF16:
+      return SizeOf<DataType::kBF16>();
+    case DataType::kF16:
+      return SizeOf<DataType::kF16>();
+    case DataType::kF32:
+      return SizeOf<DataType::kF32>();
+  }
+}
+
+// Gets a string representation of the given DataType.
+constexpr const char* ToString(DataType t) {
+  switch (t) {
+    case DataType::kI1:
+      return "I1";
+      break;
+    case DataType::kSI4:
+      return "SI4";
+      break;
+    case DataType::kSI8:
+      return "SI8";
+      break;
+    case DataType::kSI16:
+      return "SI16";
+      break;
+    case DataType::kSI32:
+      return "SI32";
+      break;
+    case DataType::kSI64:
+      return "SI64";
+      break;
+    case DataType::kBF16:
+      return "BF16";
+      break;
+    case DataType::kF16:
+      return "F16";
+      break;
+    case DataType::kF32:
+      return "F32";
+      break;
+  }
+  return "Unknown data type";
+}
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_DATA_TYPE_H_

tensorflow/lite/experimental/shlo/dispatch.h

@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_DISPATCH_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_DISPATCH_H_
+
+#define RETURN_OK_STATUS_IF_VOID(expr)                           \
+  {                                                              \
+    return [&](auto v) {                                         \
+      if constexpr (std::is_same_v<decltype(v, (expr)), void>) { \
+        (void)(expr);                                            \
+        return absl::OkStatus();                                 \
+      } else {                                                   \
+        return expr;                                             \
+      }                                                          \
+      return absl::OkStatus();                                   \
+    }(0);                                                        \
+  }
+
+#define DISPATCH_INT(name, element_type, ...)                           \
+  {                                                                     \
+    switch (element_type) {                                             \
+      case DataType::kSI4:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI4>(__VA_ARGS__)));  \
+      case DataType::kSI8:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI8>(__VA_ARGS__)));  \
+      case DataType::kSI16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI16>(__VA_ARGS__))); \
+      case DataType::kSI32:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI32>(__VA_ARGS__))); \
+      default:                                                          \
+        return absl::InvalidArgumentError("Unsupported element type");  \
+    }                                                                   \
+  }
+
+#define DISPATCH_FLOAT(name, element_type, ...)                         \
+  {                                                                     \
+    switch (element_type) {                                             \
+      case DataType::kBF16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kBF16>(__VA_ARGS__))); \
+      case DataType::kF16:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF16>(__VA_ARGS__)));  \
+      case DataType::kF32:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF32>(__VA_ARGS__)));  \
+      default:                                                          \
+        return absl::InvalidArgumentError("Unsupported element type");  \
+    }                                                                   \
+  }
+
+#define DISPATCH_INT_FLOAT(name, element_type, ...)                     \
+  {                                                                     \
+    switch (element_type) {                                             \
+      case DataType::kSI4:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI4>(__VA_ARGS__)));  \
+      case DataType::kSI8:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI8>(__VA_ARGS__)));  \
+      case DataType::kSI16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI16>(__VA_ARGS__))); \
+      case DataType::kSI32:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI32>(__VA_ARGS__))); \
+      case DataType::kBF16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kBF16>(__VA_ARGS__))); \
+      case DataType::kF16:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF16>(__VA_ARGS__)));  \
+      case DataType::kF32:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF32>(__VA_ARGS__)));  \
+      default:                                                          \
+        return absl::InvalidArgumentError("Unsupported element type");  \
+    }                                                                   \
+  }
+
+#define DISPATCH_BOOL_INT(name, element_type, ...)                      \
+  {                                                                     \
+    switch (element_type) {                                             \
+      case DataType::kI1:                                               \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kI1>(__VA_ARGS__)));   \
+      case DataType::kSI4:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI4>(__VA_ARGS__)));  \
+      case DataType::kSI8:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI8>(__VA_ARGS__)));  \
+      case DataType::kSI16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI16>(__VA_ARGS__))); \
+      case DataType::kSI32:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI32>(__VA_ARGS__))); \
+      default:                                                          \
+        return absl::InvalidArgumentError("Unsupported element type");  \
+    }                                                                   \
+  }
+
+#define DISPATCH_BOOL_INT_FLOAT(name, element_type, ...)                \
+  {                                                                     \
+    switch (element_type) {                                             \
+      case DataType::kI1:                                               \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kI1>(__VA_ARGS__)));   \
+      case DataType::kSI4:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI4>(__VA_ARGS__)));  \
+      case DataType::kSI8:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI8>(__VA_ARGS__)));  \
+      case DataType::kSI16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI16>(__VA_ARGS__))); \
+      case DataType::kSI32:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kSI32>(__VA_ARGS__))); \
+      case DataType::kBF16:                                             \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kBF16>(__VA_ARGS__))); \
+      case DataType::kF16:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF16>(__VA_ARGS__)));  \
+      case DataType::kF32:                                              \
+        RETURN_OK_STATUS_IF_VOID((name<DataType::kF32>(__VA_ARGS__)));  \
+      default:                                                          \
+        return absl::InvalidArgumentError("Unsupported element type");  \
+    }                                                                   \
+  }
+
+#define DISPATCH_QUANTIZED(name, storage_type, expressed_type, ...)          \
+  {                                                                          \
+    switch (storage_type) {                                                  \
+      case DataType::kSI4:                                                   \
+        switch (expressed_type) {                                            \
+          case DataType::kBF16:                                              \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI4, DataType::kBF16>(__VA_ARGS__)));       \
+          case DataType::kF16:                                               \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI4, DataType::kF16>(__VA_ARGS__)));        \
+          case DataType::kF32:                                               \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI4, DataType::kF32>(__VA_ARGS__)));        \
+          default:                                                           \
+            return absl::InvalidArgumentError("Unsupported expressed type"); \
+        }                                                                    \
+        break;                                                               \
+      case DataType::kSI8:                                                   \
+        switch (expressed_type) {                                            \
+          case DataType::kBF16:                                              \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI8, DataType::kBF16>(__VA_ARGS__)));       \
+          case DataType::kF16:                                               \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI8, DataType::kF16>(__VA_ARGS__)));        \
+          case DataType::kF32:                                               \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI8, DataType::kF32>(__VA_ARGS__)));        \
+          default:                                                           \
+            return absl::InvalidArgumentError("Unsupported expressed type"); \
+        }                                                                    \
+        break;                                                               \
+      case DataType::kSI16:                                                  \
+        switch (expressed_type) {                                            \
+          case DataType::kF32:                                               \
+            RETURN_OK_STATUS_IF_VOID(                                        \
+                (name<DataType::kSI16, DataType::kF32>(__VA_ARGS__)));       \
+          default:                                                           \
+            return absl::InvalidArgumentError("Unsupported expressed type"); \
+        }                                                                    \
+        break;                                                               \
+      default:                                                               \
+        return absl::InvalidArgumentError("Unsupported storage type");       \
+    }                                                                        \
+  }
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_DISPATCH_H_

tensorflow/lite/experimental/shlo/dispatch_test.cc

@@ -0,0 +1,44 @@
+/* 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/dispatch.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+
+namespace {
+
+void VoidFunction() {}
+
+TEST(DispatchTest, ReturnAbslOkIfVoidCompiles) {
+  auto f = []() -> absl::Status { RETURN_OK_STATUS_IF_VOID(VoidFunction()); };
+  EXPECT_OK(f());
+}
+
+TEST(DispatchTest, AbslOkStatusCompiles) {
+  auto f = []() -> absl::Status { RETURN_OK_STATUS_IF_VOID(absl::OkStatus()); };
+  EXPECT_OK(f());
+}
+
+TEST(DispatchTest, AbslErrorCompiles) {
+  auto f = []() -> absl::Status {
+    RETURN_OK_STATUS_IF_VOID(absl::UnknownError("error message"));
+  };
+  EXPECT_EQ(f(), absl::UnknownError("error message"));
+}
+
+}  // namespace

tensorflow/lite/experimental/shlo/f16.h

@@ -0,0 +1,245 @@
+/* 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_F16_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_F16_H_
+
+#include <cstdint>
+#include <type_traits>
+
+#include "fp16.h"  // from @FP16  // IWYU pragma: keep, used with no builtin float16
+
+// Use __FLT16_MAX__ to determine whether _Float16 is builtin
+#if defined(__FLT16_MAX__) && !SHLO_REF_EMULATE_F16
+#define SHLO_REF_HAS_BUILTIN_FLOAT16 1
+#endif
+
+namespace shlo_ref {
+
+class alignas(uint16_t) F16 {
+ public:
+  F16() = default;
+
+  template <typename T,
+            typename = std::enable_if_t<std::is_convertible_v<T, float>>>
+  // Allow implicit conversions from types convertible to float.
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  F16(T x);
+
+  // Tagged constructor to allow construction from bits
+  struct bitcast_construct_t {};
+  explicit F16(bitcast_construct_t, uint16_t bits) : bits_(bits) {}
+
+  // Allow implicit conversions to float.
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  operator float() const;
+
+  explicit operator bool() const;
+
+  F16& operator=(float x) { return *this = static_cast<F16>(x); }
+
+#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
+
+#define SHLO_REF_DEFINE_BINARY_OP(OP)                                         \
+  friend F16 operator OP(F16 x, F16 y) { return x.native_ OP y.native_; }     \
+                                                                              \
+  template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
+  friend auto operator OP(F16 x, T y) {                                       \
+    return x.native_ OP y;                                                    \
+  }                                                                           \
+                                                                              \
+  template <typename T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
+  friend auto operator OP(T x, F16 y) {                                       \
+    return x OP y.native_;                                                    \
+  }
+
+#define SHLO_REF_DEFINE_BINARY_ASSIGN_OP(OP)                               \
+  SHLO_REF_DEFINE_BINARY_OP(OP);                                           \
+  friend F16& operator OP##=(F16 & x, F16 y) {                             \
+    x.native_ OP## = y.native_;                                            \
+    return x;                                                              \
+  }                                                                        \
+                                                                           \
+  template <class T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
+  friend F16& operator OP##=(F16 & x, T y) {                               \
+    x.native_ OP## = y;                                                    \
+    return x;                                                              \
+  }
+
+#else  // !SHLO_REF_HAS_BUILTIN_FLOAT16
+
+#define SHLO_REF_DEFINE_BINARY_OP(OP)                              \
+  friend F16 operator OP(F16 x, F16 y) {                           \
+    return F16(static_cast<float>(x) OP static_cast<float>(y));    \
+  }                                                                \
+                                                                   \
+  template <typename T, typename C = std::common_type_t<F16, T>>   \
+  friend C operator OP(F16 x, T y) {                               \
+    return static_cast<C>(static_cast<C>(x) OP static_cast<C>(y)); \
+  }                                                                \
+                                                                   \
+  template <typename T, typename C = std::common_type_t<F16, T>>   \
+  friend C operator OP(T x, F16 y) {                               \
+    return static_cast<C>(static_cast<C>(x) OP static_cast<C>(y)); \
+  }
+
+#define SHLO_REF_DEFINE_BINARY_ASSIGN_OP(OP)                               \
+  SHLO_REF_DEFINE_BINARY_OP(OP);                                           \
+  friend F16& operator OP##=(F16 & x, F16 y) {                             \
+    return x = F16(static_cast<float>(x) OP static_cast<float>(y));        \
+  }                                                                        \
+                                                                           \
+  template <class T, typename = std::enable_if_t<std::is_arithmetic_v<T>>> \
+  friend F16& operator OP##=(F16 & x, T y) {                               \
+    return x = static_cast<float>(x) OP y;                                 \
+  }
+
+#endif  // SHLO_REF_HAS_BUILTIN_FLOAT16
+
+  friend F16 operator+(F16 x);
+  friend F16 operator-(F16 x);
+  friend F16& operator++(F16& x);
+  friend F16 operator++(F16& x, int);
+  friend F16& operator--(F16& x);
+  friend F16 operator--(F16& x, int);
+  SHLO_REF_DEFINE_BINARY_ASSIGN_OP(+);
+  SHLO_REF_DEFINE_BINARY_ASSIGN_OP(-);
+  SHLO_REF_DEFINE_BINARY_ASSIGN_OP(*);
+  SHLO_REF_DEFINE_BINARY_ASSIGN_OP(/);
+  SHLO_REF_DEFINE_BINARY_OP(<);
+  SHLO_REF_DEFINE_BINARY_OP(<=);
+  SHLO_REF_DEFINE_BINARY_OP(>);
+  SHLO_REF_DEFINE_BINARY_OP(>=);
+  SHLO_REF_DEFINE_BINARY_OP(==);
+  SHLO_REF_DEFINE_BINARY_OP(!=);
+#undef SHLO_REF_DEFINE_BINARY_ASSIGN_OP
+#undef SHLO_REF_DEFINE_BINARY_OP
+
+ private:
+  union {
+#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
+    _Float16 native_;
+#endif
+
+    uint16_t bits_;
+  };
+};
+
+namespace detail {
+
+template <class T, class SFINAE = void>
+struct F16CommonType {};
+
+template <class T>
+struct F16CommonType<T, std::enable_if_t<std::is_integral_v<T>>> {
+  using type = F16;
+};
+
+template <class T>
+struct F16CommonType<T, std::enable_if_t<std::is_floating_point_v<T>>> {
+  using type = T;
+};
+
+template <class T>
+struct F16CommonType<T, std::enable_if_t<!std::is_arithmetic_v<T> &&
+                                         std::is_convertible_v<T, float>>> {
+  using type = float;
+};
+
+template <class T>
+struct F16CommonType<T, std::enable_if_t<!std::is_arithmetic_v<T> &&
+                                         !std::is_convertible_v<T, float> &&
+                                         std::is_convertible_v<T, double>>> {
+  using type = double;
+};
+
+}  // namespace detail
+
+}  // namespace shlo_ref
+
+namespace std {
+
+template <>
+struct common_type<shlo_ref::F16, shlo_ref::F16> {
+  using type = shlo_ref::F16;
+};
+
+template <typename T>
+struct common_type<shlo_ref::F16, T> : shlo_ref::detail::F16CommonType<T> {};
+
+template <typename T>
+struct common_type<T, shlo_ref::F16> : shlo_ref::detail::F16CommonType<T> {};
+
+}  // namespace std
+
+namespace shlo_ref {
+static_assert(sizeof(F16) == sizeof(uint16_t));
+static_assert(alignof(F16) == alignof(uint16_t));
+
+#ifdef SHLO_REF_HAS_BUILTIN_FLOAT16
+
+template <typename T, typename _>
+F16::F16(T x) : native_(static_cast<_Float16>(static_cast<float>(x))) {}
+
+inline F16::operator float() const { return native_; }
+
+inline F16::operator bool() const { return native_; }
+
+inline F16 operator+(F16 x) { return x.native_; }
+
+inline F16 operator-(F16 x) { return -x.native_; }
+
+inline F16& operator++(F16& x) { return x += 1; }
+
+inline F16 operator++(F16& x, int) { return x.native_++; }
+
+inline F16& operator--(F16& x) { return x -= 1; }
+
+inline F16 operator--(F16& x, int) { return x.native_--; }
+
+#else  // !SHLO_REF_HAS_BUILTIN_FLOAT16
+
+template <typename T, typename _>
+inline F16::F16(T x)
+    : bits_(fp16_ieee_from_fp32_value(static_cast<float>(x))) {}
+
+inline F16::operator float() const { return fp16_ieee_to_fp32_value(bits_); }
+
+inline F16::operator bool() const { return bits_; }
+
+inline F16 operator-(F16 x) { return F16(-static_cast<float>(x)); }
+inline F16 operator+(F16 x) { return F16(static_cast<float>(x)); }
+
+inline F16& operator++(F16& x) { return x += 1; }
+
+inline F16 operator++(F16& x, int) {
+  const F16 y = x;
+  ++x;
+  return y;
+}
+
+inline F16& operator--(F16& x) { return x -= 1; }
+
+inline F16 operator--(F16& x, int) {
+  const F16 y = x;
+  --x;
+  return y;
+}
+
+#endif
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_F16_H_

tensorflow/lite/experimental/shlo/f16_test.cc

@@ -0,0 +1,188 @@
+/* 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/f16.h"
+
+#include <cstdint>
+#include <limits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+
+namespace shlo_ref {
+namespace {
+
+using ::testing::FloatNear;
+
+using RoundtripTypeList = ::testing::Types<float, double>;
+
+template <class T>
+struct RoundtripF16Test : testing::Test {};
+
+TYPED_TEST_SUITE(RoundtripF16Test, RoundtripTypeList);
+
+TYPED_TEST(RoundtripF16Test, RoundtripConversions) {
+  for (TypeParam value : {
+           -std::numeric_limits<TypeParam>::infinity(),
+           std::numeric_limits<TypeParam>::infinity(),
+           TypeParam(-1.0),
+           TypeParam(-0.5),
+           TypeParam(-0.0),
+           TypeParam(1.0),
+           TypeParam(0.5),
+           TypeParam(0.0),
+       }) {
+    EXPECT_EQ(value, static_cast<TypeParam>(static_cast<F16>(value)));
+  }
+}
+
+TEST(F16Test, Arithmetic) {
+  EXPECT_EQ(static_cast<float>(F16(2) + F16(2)), 4);
+  EXPECT_EQ(static_cast<float>(F16(2) + F16(-2)), 0);
+  EXPECT_THAT(static_cast<float>(F16(0.33333f) + F16(0.66667f)),
+              FloatNear(1.0f, 1e-3));
+  EXPECT_EQ(static_cast<float>(F16(2.0f) * F16(-5.5f)), -11.0f);
+  EXPECT_THAT(static_cast<float>(F16(1.0f) / F16(3.0f)),
+              FloatNear(0.3339f, 1e-3));
+  EXPECT_EQ(static_cast<float>(-F16(4096.0f)), -4096.0f);
+  EXPECT_EQ(static_cast<float>(-F16(-4096.0f)), 4096.0f);
+}
+
+TEST(F16Test, DefaultConstruct) { EXPECT_EQ(static_cast<float>(F16()), 0.0f); }
+
+TEST(F16Test, ImplicitConversionToFloat) {
+  EXPECT_EQ((absl::bit_cast<F16, uint16_t>(0x0000)), 0.0f);
+  EXPECT_EQ((absl::bit_cast<F16, uint16_t>(0x3C00)), 1.0f);
+}
+
+TEST(F16Test, ConstructFromArithmeticType) {
+  const F16 from_int8(static_cast<int8_t>(1));
+  EXPECT_EQ(static_cast<float>(from_int8), 1);
+  const F16 from_int16(static_cast<int16_t>(1));
+  EXPECT_EQ(static_cast<float>(from_int16), 1);
+  const F16 from_int32(static_cast<int32_t>(1));
+  EXPECT_EQ(static_cast<float>(from_int32), 1);
+  const F16 from_int64(static_cast<int64_t>(1));
+  EXPECT_EQ(static_cast<float>(from_int64), 1);
+  const F16 from_float(static_cast<float>(1));
+  EXPECT_EQ(static_cast<float>(from_float), 1);
+  const F16 from_double(static_cast<double>(1));
+  EXPECT_EQ(static_cast<float>(from_double), 1);
+}
+
+template <class T>
+T ImplicitConversion(T v) {
+  return v;
+}
+
+TEST(F16Test, ConvertToArithmeticType) {
+  const F16 ref(-1);
+  EXPECT_EQ(ImplicitConversion<int8_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int16_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int32_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int64_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<float>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<double>(ref), -1);
+}
+
+TEST(F16Test, ArithmeticOperations) {
+  // Every test relies on the equality comparisons working. We test all the 4
+  // bit integral values.
+  for (int i = -8; i < 8; ++i) {
+    for (int j = -8; j < 8; ++j) {
+      EXPECT_EQ(F16(i) == F16(j), i == j);
+      EXPECT_EQ(F16(i) != F16(j), i != j);
+      EXPECT_EQ(F16(i) > F16(j), i > j);
+      EXPECT_EQ(F16(i) >= F16(j), i >= j);
+      EXPECT_EQ(F16(i) < F16(j), i < j);
+      EXPECT_EQ(F16(i) <= F16(j), i <= j);
+    }
+  }
+  F16 val(0);
+  EXPECT_EQ(++val, 1);
+  EXPECT_EQ(val++, 1);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val--, 2);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(--val, 0);
+  EXPECT_EQ(val += F16(1), 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val *= F16(2), 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val /= F16(2), 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val -= F16(4), -3);
+  EXPECT_EQ(val, -3);
+  EXPECT_EQ(val = F16(7), 7);
+  EXPECT_EQ(val, 7);
+  EXPECT_EQ(+val, 7);
+  EXPECT_EQ(-val, -7);
+  EXPECT_EQ(static_cast<bool>(val), true);
+  EXPECT_EQ(!val, false);
+  EXPECT_EQ(val && F16(2), true);
+  EXPECT_EQ(val && F16(0), false);
+  EXPECT_EQ(val || F16(0), true);
+  EXPECT_EQ(F16(0) || F16(0), false);
+}
+
+using ArithmeticTypeList =
+    ::testing::Types<int8_t, int16_t, int32_t, int64_t, float, double>;
+
+template <class T>
+struct ArithmeticTypeF16Test : testing::Test {};
+
+TYPED_TEST_SUITE(ArithmeticTypeF16Test, ArithmeticTypeList);
+
+TYPED_TEST(ArithmeticTypeF16Test, InPlaceArithmetic) {
+  // Every test relies on the equality comparisons working. We test all the 4
+  // bit integral values.
+  for (TypeParam i = -8; i < 8; ++i) {
+    for (TypeParam j = -8; j < 8; ++j) {
+      EXPECT_EQ(F16(i) == j, i == j);
+      EXPECT_EQ(i == F16(j), i == j);
+      EXPECT_EQ(F16(i) != j, i != j);
+      EXPECT_EQ(i != F16(j), i != j);
+      EXPECT_EQ(F16(i) > j, i > j);
+      EXPECT_EQ(i > F16(j), i > j);
+      EXPECT_EQ(F16(i) >= j, i >= j);
+      EXPECT_EQ(i >= F16(j), i >= j);
+      EXPECT_EQ(F16(i) < j, i < j);
+      EXPECT_EQ(i < F16(j), i < j);
+      EXPECT_EQ(F16(i) <= j, i <= j);
+      EXPECT_EQ(i <= F16(j), i <= j);
+    }
+  }
+  const TypeParam one = TypeParam(1);
+  const TypeParam two = TypeParam(2);
+  const TypeParam four = TypeParam(4);
+  F16 val(0);
+  EXPECT_EQ(val += one, 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val *= two, 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val /= two, 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val -= four, -3);
+  EXPECT_EQ(val, -3);
+  const F16 f16_three(3);
+  EXPECT_EQ(f16_three + one, 4.);
+  EXPECT_EQ(f16_three - one, 2.);
+  EXPECT_EQ(f16_three * two, 3. * two);
+  EXPECT_EQ(f16_three / two, 3. / two);
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/i4.h

@@ -0,0 +1,416 @@
+/* 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_I4_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_I4_H_
+
+#include <cstdint>
+#include <limits>
+#include <ostream>
+#include <type_traits>
+
+namespace shlo_ref {
+
+struct I4 {
+  int8_t data = 0;
+
+  constexpr I4() = default;
+  constexpr I4(const I4&) = default;
+  constexpr I4& operator=(const I4&) = default;
+
+  template <class T, class = std::enable_if_t<std::is_convertible_v<T, int8_t>>>
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  constexpr I4(T v) : data(v) {}
+
+  template <class T, class = std::enable_if_t<std::is_convertible_v<int8_t, T>>>
+  // NOLINTNEXTLINE(google-explicit-constructor)
+  constexpr operator T() const {
+    return static_cast<T>(data);
+  }
+
+  // ++a
+  friend I4& operator++(I4& lhs) {
+    ++lhs.data;
+    return lhs;
+  }
+  // --a
+  friend I4& operator--(I4& lhs) {
+    --lhs.data;
+    return lhs;
+  }
+  // a++
+  friend I4 operator++(I4& lhs, int) {
+    I4 ret = lhs;
+    ++lhs.data;
+    return ret;
+  }
+  // a--
+  friend I4 operator--(I4& lhs, int) {
+    I4 ret = lhs;
+    --lhs.data;
+    return ret;
+  }
+  // a += b
+  friend I4& operator+=(I4& lhs, I4 rhs) {
+    lhs.data += rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator+=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data += static_cast<C>(rhs);
+    return lhs;
+  }
+  // a -= b
+  friend I4& operator-=(I4& lhs, I4 rhs) {
+    lhs.data -= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator-=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data -= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a *= b
+  friend I4& operator*=(I4& lhs, I4 rhs) {
+    lhs.data *= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator*=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data *= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a /= b
+  friend I4& operator/=(I4& lhs, I4 rhs) {
+    lhs.data /= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator/=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data /= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a %= b
+  friend I4& operator%=(I4& lhs, I4 rhs) {
+    lhs.data %= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator%=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data %= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a &= b
+  friend I4& operator&=(I4& lhs, I4 rhs) {
+    lhs.data &= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator&=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data &= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a |= b
+  friend I4& operator|=(I4& lhs, I4 rhs) {
+    lhs.data |= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator|=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data |= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a ^= b
+  friend I4& operator^=(I4& lhs, I4 rhs) {
+    lhs.data ^= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator^=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data ^= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a <<= b
+  friend I4& operator<<=(I4& lhs, I4 rhs) {
+    lhs.data <<= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator<<=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data <<= static_cast<C>(rhs);
+    return lhs;
+  }
+  // a >>= b
+  friend I4& operator>>=(I4& lhs, I4 rhs) {
+    lhs.data >>= rhs.data;
+    return lhs;
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend I4& operator>>=(I4& lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    lhs.data >>= static_cast<C>(rhs);
+    return lhs;
+  }
+  // +a
+  friend auto operator+(I4 lhs) { return +lhs.data; }
+  // -a
+  friend auto operator-(I4 lhs) { return -lhs.data; }
+  // a + b
+  friend auto operator+(I4 lhs, I4 rhs) { return lhs.data + rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator+(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data + static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator+(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) + rhs.data;
+  }
+  // a - b
+  friend auto operator-(I4 lhs, I4 rhs) { return lhs.data - rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator-(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data - static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator-(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) - rhs.data;
+  }
+  // a * b
+  friend auto operator*(I4 lhs, I4 rhs) { return lhs.data * rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator*(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data * static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator*(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) * rhs.data;
+  }
+  // a / b
+  friend auto operator/(I4 lhs, I4 rhs) { return lhs.data / rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator/(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data / static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator/(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) / rhs.data;
+  }
+  // a % b
+  friend auto operator%(I4 lhs, I4 rhs) { return lhs.data % rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator%(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data % static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator%(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) % rhs.data;
+  }
+  // ~a
+  friend auto operator~(I4 lhs) { return ~lhs.data; }
+  // a & b
+  friend auto operator&(I4 lhs, I4 rhs) { return lhs.data & rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator&(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data & static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator&(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) & rhs.data;
+  }
+  // a | b
+  friend auto operator|(I4 lhs, I4 rhs) { return lhs.data | rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator|(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data | static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator|(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) | rhs.data;
+  }
+  // a ^ b
+  friend auto operator^(I4 lhs, I4 rhs) { return lhs.data ^ rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator^(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data ^ static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator^(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) ^ rhs.data;
+  }
+  // a << b
+  friend auto operator<<(I4 lhs, I4 rhs) { return lhs.data << rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator<<(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data << static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator<<(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) << rhs.data;
+  }
+  // a >> b
+  friend auto operator>>(I4 lhs, I4 rhs) { return lhs.data >> rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator>>(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data >> static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_integral_v<T>>>
+  friend auto operator>>(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) >> rhs.data;
+  }
+  // !a
+  friend bool operator!(I4 v) { return !v.data; }
+  // a && b
+  friend auto operator&&(I4 lhs, I4 rhs) { return lhs.data && rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator&&(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data && static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator&&(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) && rhs.data;
+  }
+  // a || b
+  friend auto operator||(I4 lhs, I4 rhs) { return lhs.data || rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator||(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data || static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend auto operator||(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) || rhs.data;
+  }
+  // a == b
+  friend bool operator==(I4 lhs, I4 rhs) { return lhs.data == rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator==(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data == static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator==(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) == rhs.data;
+  }
+  // a != b
+  friend bool operator!=(I4 lhs, I4 rhs) { return lhs.data != rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator!=(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data != static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator!=(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) != rhs.data;
+  }
+  // a < b
+  friend bool operator<(I4 lhs, I4 rhs) { return lhs.data < rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator<(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data < static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator<(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) < rhs.data;
+  }
+  // a > b
+  friend bool operator>(I4 lhs, I4 rhs) { return lhs.data > rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator>(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data > static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator>(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) > rhs.data;
+  }
+  // a <= b
+  friend bool operator<=(I4 lhs, I4 rhs) { return lhs.data <= rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator<=(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data <= static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator<=(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) <= rhs.data;
+  }
+  // a >= b
+  friend bool operator>=(I4 lhs, I4 rhs) { return lhs.data >= rhs.data; }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator>=(I4 lhs, T rhs) {
+    using C = std::common_type_t<T, int>;
+    return lhs.data >= static_cast<C>(rhs);
+  }
+  template <class T, class = std::enable_if_t<std::is_arithmetic_v<T>>>
+  friend bool operator>=(T lhs, I4 rhs) {
+    using C = std::common_type_t<T, int>;
+    return static_cast<C>(lhs) >= rhs.data;
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, I4 v) { return os << +v; }
+};
+
+}  // namespace shlo_ref
+
+namespace std {
+
+template <>
+struct numeric_limits<shlo_ref::I4> : std::numeric_limits<int8_t> {
+  static constexpr shlo_ref::I4 min() noexcept { return shlo_ref::I4(-8); }
+  static constexpr shlo_ref::I4 lowest() noexcept { return min(); }
+  static constexpr shlo_ref::I4 max() noexcept { return shlo_ref::I4(7); }
+};
+
+}  // namespace std
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_I4_H_

tensorflow/lite/experimental/shlo/i4_test.cc

@@ -0,0 +1,198 @@
+/* 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/i4.h"
+
+#include <cstdint>
+
+#include <gtest/gtest.h>
+
+namespace shlo_ref {
+namespace {
+
+TEST(I4Test, ConstructFromArithmeticType) {
+  const I4 from_int8(static_cast<int8_t>(1));
+  EXPECT_EQ(from_int8.data, 1);
+  const I4 from_int16(static_cast<int16_t>(1));
+  EXPECT_EQ(from_int16.data, 1);
+  const I4 from_int32(static_cast<int32_t>(1));
+  EXPECT_EQ(from_int32.data, 1);
+  const I4 from_int64(static_cast<int64_t>(1));
+  EXPECT_EQ(from_int64.data, 1);
+  const I4 from_float(static_cast<float>(1));
+  EXPECT_EQ(from_float.data, 1);
+  const I4 from_double(static_cast<double>(1));
+  EXPECT_EQ(from_double.data, 1);
+}
+
+template <class T>
+T ImplicitConversion(T v) {
+  return v;
+}
+
+TEST(I4Test, ConvertToArithmeticType) {
+  const I4 ref(-1);
+  EXPECT_EQ(ImplicitConversion<int8_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int16_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int32_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<int64_t>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<float>(ref), -1);
+  EXPECT_EQ(ImplicitConversion<double>(ref), -1);
+}
+
+TEST(I4Test, Arithmetic) {
+  // Every test relies on the equality comparisons working. We test all the 4
+  // bit integral values.
+  for (int i = -8; i < 8; ++i) {
+    for (int j = -8; j < 8; ++j) {
+      EXPECT_EQ(I4(i) == I4(j), i == j);
+      EXPECT_EQ(I4(i) != I4(j), i != j);
+      EXPECT_EQ(I4(i) > I4(j), i > j);
+      EXPECT_EQ(I4(i) >= I4(j), i >= j);
+      EXPECT_EQ(I4(i) < I4(j), i < j);
+      EXPECT_EQ(I4(i) <= I4(j), i <= j);
+    }
+  }
+  I4 val(0);
+  EXPECT_EQ(++val, 1);
+  EXPECT_EQ(val++, 1);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val--, 2);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(--val, 0);
+  EXPECT_EQ(val += I4(1), 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val *= I4(2), 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val /= I4(2), 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val -= I4(4), -3);
+  EXPECT_EQ(val, -3);
+  EXPECT_EQ(val %= I4(2), -1);
+  EXPECT_EQ(val, -1);
+  EXPECT_EQ(val = I4(7), 7);
+  EXPECT_EQ(val, 7);
+  EXPECT_EQ(val &= I4(2), 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val |= I4(1), 3);
+  EXPECT_EQ(val, 3);
+  EXPECT_EQ(val ^= I4(7), 4);
+  EXPECT_EQ(val, 4);
+  EXPECT_EQ(val >>= I4(1), 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val <<= I4(1), 4);
+  EXPECT_EQ(val, 4);
+  EXPECT_EQ(val >>= I4(1), 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val <<= I4(1), 4);
+  EXPECT_EQ(val, 4);
+  EXPECT_EQ(+val, 4);
+  EXPECT_EQ(-val, -4);
+  EXPECT_EQ(!val, false);
+  EXPECT_EQ(~val, ~4);
+  EXPECT_EQ(val && I4(2), true);
+  EXPECT_EQ(val && I4(0), false);
+  EXPECT_EQ(val || I4(0), true);
+  EXPECT_EQ(I4(0) || I4(0), false);
+}
+
+using IntegralTypeList = ::testing::Types<int8_t, int16_t, int32_t, int64_t>;
+
+using ArithmeticTypeList =
+    ::testing::Types<int8_t, int16_t, int32_t, int64_t, float, double>;
+
+template <class T>
+struct ArithmeticTypeI4Test : testing::Test {};
+
+TYPED_TEST_SUITE(ArithmeticTypeI4Test, ArithmeticTypeList);
+
+TYPED_TEST(ArithmeticTypeI4Test, Arithmetic) {
+  // Every test relies on the equality comparisons working. We test all the 4
+  // bit integral values.
+  for (TypeParam i = -8; i < 8; ++i) {
+    for (TypeParam j = -8; j < 8; ++j) {
+      EXPECT_EQ(I4(i) == j, i == j);
+      EXPECT_EQ(i == I4(j), i == j);
+      EXPECT_EQ(I4(i) != j, i != j);
+      EXPECT_EQ(i != I4(j), i != j);
+      EXPECT_EQ(I4(i) > j, i > j);
+      EXPECT_EQ(i > I4(j), i > j);
+      EXPECT_EQ(I4(i) >= j, i >= j);
+      EXPECT_EQ(i >= I4(j), i >= j);
+      EXPECT_EQ(I4(i) < j, i < j);
+      EXPECT_EQ(i < I4(j), i < j);
+      EXPECT_EQ(I4(i) <= j, i <= j);
+      EXPECT_EQ(i <= I4(j), i <= j);
+    }
+  }
+  I4 val(0);
+  const TypeParam one = TypeParam(1);
+  const TypeParam two = TypeParam(2);
+  const TypeParam three = TypeParam(3);
+  const TypeParam four = TypeParam(4);
+  EXPECT_EQ(val += one, 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val *= two, 2);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val /= two, 1);
+  EXPECT_EQ(val, 1);
+  EXPECT_EQ(val -= four, -3);
+  EXPECT_EQ(val, -3);
+  const I4 i4_three(3);
+  EXPECT_EQ(i4_three + one, four);
+  EXPECT_EQ(i4_three - one, two);
+  EXPECT_EQ(i4_three * two, three * two);
+  EXPECT_EQ(i4_three / two, three / two);
+}
+
+template <class T>
+struct IntegralTypeI4Test : testing::Test {};
+
+TYPED_TEST_SUITE(IntegralTypeI4Test, IntegralTypeList);
+
+TYPED_TEST(IntegralTypeI4Test, Arithmetic) {
+  const TypeParam minus_one = TypeParam(-1);
+  const TypeParam one = TypeParam(1);
+  const TypeParam two = TypeParam(2);
+  const TypeParam three = TypeParam(3);
+  const TypeParam four = TypeParam(4);
+  const TypeParam six = TypeParam(6);
+  const TypeParam seven = TypeParam(7);
+  const I4 i4_three(3);
+  EXPECT_EQ(i4_three % two, one);
+  EXPECT_EQ(i4_three & two, two);
+  EXPECT_EQ(i4_three | four, seven);
+  EXPECT_EQ(i4_three ^ four, seven);
+  EXPECT_EQ(i4_three << one, six);
+  EXPECT_EQ(i4_three >> one, one);
+  I4 val(-3);
+  EXPECT_EQ(val %= two, minus_one);
+  EXPECT_EQ(val, -1);
+  EXPECT_EQ(val = I4(7), seven);
+  EXPECT_EQ(val, 7);
+  EXPECT_EQ(val &= two, two);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val |= one, three);
+  EXPECT_EQ(val, 3);
+  EXPECT_EQ(val ^= seven, four);
+  EXPECT_EQ(val, 4);
+  EXPECT_EQ(val >>= one, two);
+  EXPECT_EQ(val, 2);
+  EXPECT_EQ(val <<= one, four);
+  EXPECT_EQ(val, 4);
+}
+
+}  // namespace
+}  // namespace shlo_ref

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_

tensorflow/lite/experimental/shlo/ops/abs.cc

@@ -0,0 +1,59 @@
+/* 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/ops/abs.h"
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Abs {
+  template <class T>
+  T operator()(const T& val) {
+    return val < static_cast<T>(0) ? static_cast<T>(-val) : val;
+  }
+};
+
+AbsOp Create(typename AbsOp::Attributes) { return AbsOp{}; }
+
+absl::Status Prepare(AbsOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(CheckCtx("abs"), input,
+                                               IsSignedIntTensor, IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("abs"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(AbsOp& op, const Tensor& input, Tensor& output) {
+  Abs abs;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), abs, input,
+        output)
+  } else if (IsSignedIntTensor(input) || IsFloatTensor(input)) {
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       input.tensor_element_type(), abs, input, output);
+  }
+  return absl::FailedPreconditionError("Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/abs.h

@@ -0,0 +1,33 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_ABS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_ABS_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct AbsOp {
+  struct Attributes {};
+};
+
+AbsOp Create(AbsOp::Attributes);
+absl::Status Prepare(AbsOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(AbsOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_ABS_H_

tensorflow/lite/experimental/shlo/ops/abs_test.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 "tensorflow/lite/experimental/shlo/ops/abs.h"
+
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<AbsOp> {
+  static std::string Get() { return "Abs"; }
+};
+
+namespace {
+
+constexpr struct AbsRef {
+  template <class T>
+  T operator()(T v) const {
+    return v < static_cast<T>(0) ? static_cast<T>(-v) : v;
+  }
+} abs_ref;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Abs, UnaryElementwiseOpShapePropagationTest,
+                               AbsOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Abs, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<AbsOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<AbsOp, ConcatTypes<BoolTestType, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Abs, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct AbsTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(AbsTest, ArithmeticTestTypes, TestParamNames);
+
+TYPED_TEST(AbsTest, ArithmeticTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), abs_ref);
+
+  auto op = Create(AbsOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+template <class T>
+struct QuantizedAbsTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedAbsTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedAbsTest, QuantizedPerTensor) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = abs_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(AbsOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/and.cc

@@ -0,0 +1,63 @@
+/* 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/ops/and.h"
+
+#include <functional>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+template <DataType>
+struct And : std::bit_and<void> {};
+
+template <>
+struct And<DataType::kI1> : std::logical_and<void> {};
+
+AndOp Create(AndOp::Attributes) { return {}; }
+
+absl::Status Prepare(AndOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(lhs.shape(), rhs.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("and"), lhs, IsBoolTensor, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("and"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("and"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(AndOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  if (IsIntTensor(lhs)) {
+    // Note: all the integer types share the same implementation.
+    And<DataType::kSI32> and_func;
+    DISPATCH_INT(detail::EvaluateNoQuantization, lhs.tensor_element_type(),
+                 and_func, lhs, rhs, output);
+  } else if (IsBoolTensor(lhs)) {
+    And<DataType::kI1> and_func;
+    detail::EvaluateNoQuantization<DataType::kI1>(and_func, lhs, rhs, output);
+    return absl::OkStatus();
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.and: Unsupported tensor type in Evaluate.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/and.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_OPS_AND_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_AND_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct AndOp {
+  struct Attributes {};
+};
+
+AndOp Create(AndOp::Attributes);
+absl::Status Prepare(AndOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output);
+absl::Status Evaluate(AndOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_AND_H_

tensorflow/lite/experimental/shlo/ops/and_test.cc

@@ -0,0 +1,107 @@
+/* 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/ops/and.h"
+
+#include <functional>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::FloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<AndOp> {
+  static std::string Get() { return "And"; }
+};
+
+template <DataType>
+struct And : std::bit_and<void> {};
+
+template <>
+struct And<DataType::kI1> : std::logical_and<void> {};
+
+template <>
+struct SupportedOpDataType<AndOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(And, BinaryElementwiseOpShapePropagationTest,
+                               AndOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    AndOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    And, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<AndOp, ConcatTypes<FloatTestTypes, PerTensorQuantizedTestTypes,
+                                   PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(And, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, IntTestTypes>;
+
+template <class T>
+struct AndTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(AndTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(AndTest, ArithmeticTestTypesTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-50, /*max=*/50);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/1, /*max=*/5);
+  Vector<StorageT> output_data(shape.NumElements());
+  Tensor lhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(lhs_data, rhs_data, expected_data.begin(),
+                    And<TypeParam::kStorage>());
+
+  auto op = Create(AndOp::Attributes{});
+  ASSERT_OK(Prepare(op, lhs_tensor, rhs_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, lhs_tensor, rhs_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/benchmark_util.h

@@ -0,0 +1,61 @@
+/* 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_OPS_BENCHMARK_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BENCHMARK_UTIL_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <random>
+#include <type_traits>
+#include <vector>
+
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+
+namespace shlo_ref {
+
+// Converts the given number of Kibibytes to the equivalent number of bytes.
+// This is useful for specifying test input sizes as `KiB(8)`.
+constexpr size_t KiB(size_t kibibytes) { return kibibytes * 1024; }
+
+template <DataType data_type, typename T = StorageType<data_type>>
+std::vector<T> GenerateRandomVector(size_t size) {
+  std::vector<T> data(size);
+  if constexpr (std::is_integral_v<T>) {
+    static std::uniform_int_distribution<T> distribution(
+        Storage<data_type>::kMinValue, Storage<data_type>::kMaxValue);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else if constexpr (std::is_floating_point_v<T>) {
+    static std::uniform_real_distribution<T> distribution(-1.0f, 1.0f);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  } else {
+    static_assert(std::is_same_v<T, BF16> || std::is_same_v<T, F16>);
+    static std::uniform_real_distribution<float> distribution(-1.0f, 1.0f);
+    static std::default_random_engine generator;
+    std::generate(data.begin(), data.end(),
+                  [&]() { return distribution(generator); });
+  }
+  return data;
+}
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BENCHMARK_UTIL_H_

tensorflow/lite/experimental/shlo/ops/binary_elementwise.h

@@ -0,0 +1,81 @@
+/* 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_OPS_BINARY_ELEMENTWISE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BINARY_ELEMENTWISE_H_
+
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+namespace detail {
+
+template <DataType storage_type, DataType expressed_type, typename F>
+void DequantizeOpQuantizePerTensor(F&& func, const Tensor& lhs,
+                                   const Tensor& rhs, Tensor& output) {
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+  const DimensionSize num_elements = lhs.NumElements();
+  const StorageT lhs_zero_point =
+      lhs.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT lhs_scale =
+      lhs.quantized_per_tensor_element_type().ScaleAs<expressed_type>();
+  const StorageT rhs_zero_point =
+      rhs.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT rhs_scale =
+      rhs.quantized_per_tensor_element_type().ScaleAs<expressed_type>();
+  const StorageT output_zero_point =
+      output.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT output_scale =
+      output.quantized_per_tensor_element_type().ScaleAs<expressed_type>();
+  const StorageT* lhs_data = lhs.GetDataAs<storage_type>();
+  const StorageT* rhs_data = rhs.GetDataAs<storage_type>();
+  StorageT* output_data = output.GetDataAs<storage_type>();
+  const ExpressedT inv_scale = static_cast<ExpressedT>(1) / output_scale;
+  for (DimensionSize i = 0; i < num_elements;
+       ++i, ++lhs_data, ++rhs_data, ++output_data) {
+    const ExpressedT dequantized_lhs =
+        Dequantize(*lhs_data, lhs_zero_point, lhs_scale);
+    const ExpressedT dequantized_rhs =
+        Dequantize(*rhs_data, rhs_zero_point, rhs_scale);
+    const ExpressedT dequantized_res = func(dequantized_lhs, dequantized_rhs);
+    *output_data = Quantize<storage_type, expressed_type>(
+        dequantized_res, output_zero_point, inv_scale);
+  }
+}
+
+template <DataType data_type, class F>
+void EvaluateNoQuantization(F&& func, const Tensor& lhs, const Tensor& rhs,
+                            Tensor& output) {
+  using T = StorageType<data_type>;
+  const T* lhs_data = lhs.GetDataAs<data_type>();
+  const T* rhs_data = rhs.GetDataAs<data_type>();
+  T* output_data = output.GetDataAs<data_type>();
+  const DimensionSize num_elements = lhs.NumElements();
+  for (DimensionSize i = 0; i < num_elements;
+       ++i, ++output_data, ++lhs_data, ++rhs_data) {
+    *output_data = static_cast<T>(func(*lhs_data, *rhs_data));
+  }
+}
+
+}  // namespace detail
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BINARY_ELEMENTWISE_H_

tensorflow/lite/experimental/shlo/ops/binary_elementwise_test.cc

@@ -0,0 +1,138 @@
+/* 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/ops/binary_elementwise.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+
+namespace shlo_ref {
+namespace {
+
+struct TestOp {
+  template <typename T>
+  T operator()(const T& lhs, const T& rhs) {
+    return lhs + rhs;
+  }
+};
+
+template <class T>
+struct EvaluateNoQuantizationTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(EvaluateNoQuantizationTest, ArithmeticTestTypes,
+                 TestParamNames);
+
+TYPED_TEST(EvaluateNoQuantizationTest, ArithmeticTensorsWithTestOp) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-5, /*max=*/5);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-5, /*max=*/5);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor lhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(lhs_data, rhs_data, expected_data.begin(), TestOp());
+
+  detail::EvaluateNoQuantization<TypeParam::kStorage>(
+      TestOp(), lhs_tensor, rhs_tensor, output_tensor);
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+template <class T>
+struct DequantizeOpQuantizePerTensor : ::testing::Test {};
+
+TYPED_TEST_SUITE(DequantizeOpQuantizePerTensor, QuantizedTestTypes,
+                 TestParamNames);
+
+TYPED_TEST(DequantizeOpQuantizePerTensor, QuantizedPerTensorWithTestOp) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-5, /*max=*/5);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-5, /*max=*/5);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT lhs_scale = static_cast<ExpressedT>(1.3);
+  const StorageT lhs_zero_point = static_cast<StorageT>(4);
+  const ExpressedT rhs_scale = static_cast<ExpressedT>(1.2);
+  const StorageT rhs_zero_point = static_cast<StorageT>(5);
+  const ExpressedT output_scale = static_cast<ExpressedT>(1.5);
+  const StorageT output_zero_point = static_cast<StorageT>(3);
+  Tensor lhs_tensor{.type =
+                        QuantizedPerTensorTensorType{
+                            .shape = shape,
+                            .element_type = QuantizedElementTypePerTensor(
+                                TypeParam::kStorage, lhs_zero_point,
+                                TypeParam::kExpressed, lhs_scale)},
+                    .data = lhs_data.data()};
+  Tensor rhs_tensor{.type =
+                        QuantizedPerTensorTensorType{
+                            .shape = shape,
+                            .element_type = QuantizedElementTypePerTensor(
+                                TypeParam::kStorage, rhs_zero_point,
+                                TypeParam::kExpressed, rhs_scale)},
+                    .data = rhs_data.data()};
+  Tensor output_tensor{.type =
+                           QuantizedPerTensorTensorType{
+                               .shape = shape,
+                               .element_type = QuantizedElementTypePerTensor(
+                                   TypeParam::kStorage, output_zero_point,
+                                   TypeParam::kExpressed, output_scale)},
+                       .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      lhs_data, rhs_data, expected_data.begin(),
+      [lhs_zero_point, lhs_scale, rhs_zero_point, rhs_scale, output_zero_point,
+       output_scale](auto lhs, auto rhs) {
+        const ExpressedT dequantized_lhs =
+            Dequantize(lhs, lhs_zero_point, lhs_scale);
+        const ExpressedT dequantized_rhs =
+            Dequantize(rhs, rhs_zero_point, rhs_scale);
+        const ExpressedT dequantized_res =
+            TestOp()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, output_zero_point,
+            static_cast<ExpressedT>(1.) / output_scale);
+      });
+
+  detail::DequantizeOpQuantizePerTensor<TypeParam::kStorage,
+                                        TypeParam::kExpressed>(
+      TestOp(), lhs_tensor, rhs_tensor, output_tensor);
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/binary_elementwise_test_util.h

@@ -0,0 +1,201 @@
+/* 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_OPS_BINARY_ELEMENTWISE_TEST_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BINARY_ELEMENTWISE_TEST_UTIL_H_
+
+#include <tuple>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+template <class Op, class List>
+struct OpTuple;
+
+template <class Op, class... Ts>
+struct OpTuple<Op, ::testing::Types<Ts...>> {
+  using Type = std::tuple<Op, Ts...>;
+};
+
+template <class Op>
+struct OpTupleFactory {
+  template <class T>
+  using WithOp = typename OpTuple<Op, T>::Type;
+};
+
+template <class Op, class SupportedTypes>
+using BinaryElementwiseBaselineConstraintTypes =
+    MapTypes<OpTupleFactory<Op>::template WithOp,
+             FilterTypes<NegatePred<SameTypes>::template Predicate,
+                         CrossProductTypes<SupportedTypes, SupportedTypes,
+                                           SupportedTypes>>>;
+
+using BaselineConstraintIntTypes = ::testing::Types<TestParam<DataType::kSI32>>;
+
+using BaselineConstraintFloatTypes =
+    ::testing::Types<TestParam<DataType::kF32>>;
+
+using BaselineConstraintQuantizedPerTensorTypes =
+    ::testing::Types<PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+                     PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>;
+
+template <class Op>
+class BinaryElementwiseOpShapePropagationTest : public ::testing::Test {
+ protected:
+  void SetRhsShape(Shape shape) { rhs_tensor_.shape() = std::move(shape); }
+  void SetOutputShape(Shape shape) {
+    output_tensor_.shape() = std::move(shape);
+  }
+  bool LhsAndOutputShapesAreEqual() const {
+    return lhs_tensor_.shape() == output_tensor_.shape();
+  }
+
+  Op op_ = Create(SupportedOpAttributes<Op>::Get());
+  Tensor lhs_tensor_ = {
+      .type = TensorType{.shape = Shape({2, 3, 4}),
+                         .element_type = SupportedOpDataType<Op>::kStorageType},
+      .data = nullptr};
+  Tensor rhs_tensor_ = {
+      .type = TensorType{.shape = Shape({2, 3, 4}),
+                         .element_type = SupportedOpDataType<Op>::kStorageType},
+      .data = nullptr};
+  Tensor output_tensor_ = {
+      .type = TensorType{.shape = Shape(),
+                         .element_type =
+                             SupportedOpOutputDataType<Op>::kStorageType},
+      .data = nullptr};
+};
+
+TYPED_TEST_SUITE_P(BinaryElementwiseOpShapePropagationTest);
+
+TYPED_TEST_P(BinaryElementwiseOpShapePropagationTest, ShapePropagationWorks) {
+  ASSERT_TRUE(this->output_tensor_.shape().empty());
+  EXPECT_OK(Prepare(this->op_, this->lhs_tensor_, this->rhs_tensor_,
+                    this->output_tensor_));
+  EXPECT_THAT(this->output_tensor_.shape(),
+              ::testing::ElementsAreArray(this->lhs_tensor_.shape()));
+}
+
+TYPED_TEST_P(BinaryElementwiseOpShapePropagationTest,
+             SmallerOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3}));
+  ASSERT_FALSE(this->LhsAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->lhs_tensor_, this->rhs_tensor_,
+              this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shapes."));
+}
+
+TYPED_TEST_P(BinaryElementwiseOpShapePropagationTest,
+             BiggerOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3, 4, 5}));
+  ASSERT_FALSE(this->LhsAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->lhs_tensor_, this->rhs_tensor_,
+              this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shapes."));
+}
+
+TYPED_TEST_P(BinaryElementwiseOpShapePropagationTest,
+             IncompatibleOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3, 5}));
+  ASSERT_FALSE(this->LhsAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->lhs_tensor_, this->rhs_tensor_,
+              this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shapes."));
+}
+
+REGISTER_TYPED_TEST_SUITE_P(BinaryElementwiseOpShapePropagationTest,
+                            ShapePropagationWorks,
+                            SmallerOutputShapeRaisesAnError,
+                            BiggerOutputShapeRaisesAnError,
+                            IncompatibleOutputShapeRaisesAnError);
+
+// Tests that the baseline element type of the input and output tensors is the
+// same.
+template <class T>
+class BinaryElementwiseSameBaselineElementTypeConstraintTest
+    : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(BinaryElementwiseSameBaselineElementTypeConstraintTest);
+
+TYPED_TEST_P(BinaryElementwiseSameBaselineElementTypeConstraintTest,
+             DifferentInputOutputStorageTypesRaiseAnError) {
+  using Op = std::tuple_element_t<0, TypeParam>;
+  using LhsTypeDesc = std::tuple_element_t<1, TypeParam>;
+  using RhsTypeDesc = std::tuple_element_t<2, TypeParam>;
+  using ResultTypeDesc = std::tuple_element_t<3, TypeParam>;
+  const Shape shape({2, 3, 4});
+  Tensor lhs_tensor{.type = TensorTypeFor(LhsTypeDesc{}, shape),
+                    .data = nullptr};
+  Tensor rhs_tensor{.type = TensorTypeFor(RhsTypeDesc{}, shape),
+                    .data = nullptr};
+  Tensor output_tensor{.type = TensorTypeFor(ResultTypeDesc{}, shape),
+                       .data = nullptr};
+  auto op = Create(typename Op::Attributes{});
+  const absl::Status status =
+      Prepare(op, lhs_tensor, rhs_tensor, output_tensor);
+  EXPECT_THAT(status, shlo_ref::testing::StatusIs(
+                          absl::StatusCode::kFailedPrecondition));
+  EXPECT_THAT(
+      status.message(),
+      ::testing::ContainsRegex(
+          "stablehlo.[_a-z]+: baseline type constraint is not satisfied"));
+}
+
+REGISTER_TYPED_TEST_SUITE_P(
+    BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    DifferentInputOutputStorageTypesRaiseAnError);
+
+// Tests that unsupported types are detected during when `Prepare` is called.
+template <class T>
+class BinaryElementwiseUnsupportedTypeTest : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(BinaryElementwiseUnsupportedTypeTest);
+
+TYPED_TEST_P(BinaryElementwiseUnsupportedTypeTest, PrepareRaisesAnError) {
+  using Op = std::tuple_element_t<0, TypeParam>;
+  using TypeDesc = std::tuple_element_t<1, TypeParam>;
+  Tensor input_tensor{.type = TensorTypeFor(TypeDesc{}, Shape({2, 3, 4})),
+                      .data = nullptr};
+  Tensor output_tensor = input_tensor;
+  auto op = Create(typename Op::Attributes{});
+  const absl::Status status =
+      Prepare(op, input_tensor, input_tensor, output_tensor);
+  EXPECT_THAT(status, shlo_ref::testing::StatusIs(
+                          absl::StatusCode::kFailedPrecondition));
+  EXPECT_THAT(status.message(),
+              ::testing::HasSubstr("Unsupported tensor type"));
+}
+
+REGISTER_TYPED_TEST_SUITE_P(BinaryElementwiseUnsupportedTypeTest,
+                            PrepareRaisesAnError);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_BINARY_ELEMENTWISE_TEST_UTIL_H_

tensorflow/lite/experimental/shlo/ops/cbrt.cc

@@ -0,0 +1,73 @@
+/* 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/ops/cbrt.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Cbrt {
+  template <class T>
+  T operator()(T v) const {
+    return std::cbrt(v);
+  }
+};
+
+template <>
+F16 Cbrt::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Cbrt::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+CbrtOp Create(CbrtOp::Attributes) { return {}; }
+
+absl::Status Prepare(CbrtOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("cbrt"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("cbrt"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(CbrtOp& op, const Tensor& input, Tensor& output) {
+  Cbrt cbrt;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), cbrt, input,
+        output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   cbrt, input, output);
+  }
+  return absl::FailedPreconditionError("Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/cbrt.h

@@ -0,0 +1,34 @@
+/* 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_OPS_CBRT_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_CBRT_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CbrtOp {
+  struct Attributes {};
+};
+
+CbrtOp Create(CbrtOp::Attributes);
+absl::Status Prepare(CbrtOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(CbrtOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_CBRT_H_

tensorflow/lite/experimental/shlo/ops/cbrt_test.cc

@@ -0,0 +1,146 @@
+/* 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/ops/cbrt.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<CbrtOp> {
+  static std::string Get() { return "Cbrt"; }
+};
+
+namespace {
+
+struct Cbrt {
+  template <class T>
+  T operator()(T v) const {
+    return std::cbrt(v);
+  }
+} cbrt_ref;
+
+template <>
+F16 Cbrt::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Cbrt::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Cbrt, UnaryElementwiseOpShapePropagationTest,
+                               CbrtOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Cbrt, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<CbrtOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    CbrtOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Cbrt, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct CbrtTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(CbrtTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(CbrtTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), cbrt_ref);
+
+  auto op = Create(CbrtOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedCbrtTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedCbrtTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedCbrtTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = cbrt_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(CbrtOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/ceil.cc

@@ -0,0 +1,73 @@
+/* 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/ops/ceil.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Ceil {
+  template <class T>
+  T operator()(T v) const {
+    return std::ceil(v);
+  }
+};
+
+template <>
+F16 Ceil::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Ceil::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+CeilOp Create(CeilOp::Attributes) { return {}; }
+
+absl::Status Prepare(CeilOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("ceil"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("ceil"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(CeilOp& op, const Tensor& input, Tensor& output) {
+  Ceil ceil;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), ceil, input,
+        output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   ceil, input, output);
+  }
+  return absl::FailedPreconditionError("Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/ceil.h

@@ -0,0 +1,34 @@
+/* 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_OPS_CEIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_CEIL_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CeilOp {
+  struct Attributes {};
+};
+
+CeilOp Create(CeilOp::Attributes);
+absl::Status Prepare(CeilOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(CeilOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_CEIL_H_

tensorflow/lite/experimental/shlo/ops/ceil_test.cc

@@ -0,0 +1,146 @@
+/* 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/ops/ceil.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<CeilOp> {
+  static std::string Get() { return "Ceil"; }
+};
+
+namespace {
+
+struct Ceil {
+  template <class T>
+  T operator()(T v) const {
+    return std::ceil(v);
+  }
+} ceil_ref;
+
+template <>
+F16 Ceil::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Ceil::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Ceil, UnaryElementwiseOpShapePropagationTest,
+                               CeilOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Ceil, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<CeilOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    CeilOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Ceil, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct CeilTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(CeilTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(CeilTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), ceil_ref);
+
+  auto op = Create(CeilOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedCeilTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedCeilTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedCeilTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = ceil_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(CeilOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/compare.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 "tensorflow/lite/experimental/shlo/ops/compare.h"
+
+#include <functional>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+namespace {
+
+template <DataType storage_type, DataType expressed_type, typename F>
+void DequantizeCompare(F&& func, const Tensor& lhs, const Tensor& rhs,
+                       Tensor& output) {
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+  const DimensionSize num_elements = lhs.NumElements();
+  const StorageT lhs_zero_point =
+      lhs.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT lhs_scale =
+      lhs.quantized_per_tensor_element_type().ScaleAs<expressed_type>();
+  const StorageT rhs_zero_point =
+      rhs.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT rhs_scale =
+      rhs.quantized_per_tensor_element_type().ScaleAs<expressed_type>();
+  const StorageT* lhs_data = lhs.GetDataAs<storage_type>();
+  const StorageT* rhs_data = rhs.GetDataAs<storage_type>();
+  bool* output_data = output.GetDataAs<DataType::kI1>();
+  for (DimensionSize i = 0; i < num_elements;
+       ++i, ++lhs_data, ++rhs_data, ++output_data) {
+    const ExpressedT dequantized_lhs =
+        Dequantize(*lhs_data, lhs_zero_point, lhs_scale);
+    const ExpressedT dequantized_rhs =
+        Dequantize(*rhs_data, rhs_zero_point, rhs_scale);
+    *output_data = func(dequantized_lhs, dequantized_rhs);
+  }
+}
+
+}  // namespace
+
+CompareOp Create(CompareOp::Attributes attributes) {
+  return {.attributes = attributes};
+}
+
+absl::Status Prepare(CompareOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("compare"), lhs, IsBoolTensor, IsIntTensor,
+                          IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("compare"), output, IsBoolTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("compare"), lhs, rhs));
+  SHLO_REF_RETURN_ON_ERROR(Propagate(lhs.shape(), rhs.shape(), output.shape()));
+  return absl::OkStatus();
+}
+
+// Huge body because of the type dispatch.
+// NOLINTNEXTLINE(google-readability-function-size)
+absl::Status Evaluate(CompareOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+#define SHLO_REF_COMPARISON_DIRECTION_CASE(DIRECTION, COMPARISON_OP)         \
+  case CompareOp::ComparisonDirection::DIRECTION: {                          \
+    if (IsBoolTensor(lhs) || IsIntTensor(lhs) || IsFloatTensor(lhs)) {       \
+      DISPATCH_BOOL_INT_FLOAT(detail::EvaluateNoQuantization,                \
+                              lhs.tensor_element_type(), COMPARISON_OP, lhs, \
+                              rhs, output);                                  \
+    } else if (IsQuantizedPerTensorTensor(lhs)) {                            \
+      DISPATCH_QUANTIZED(                                                    \
+          DequantizeCompare,                                                 \
+          lhs.quantized_per_tensor_element_type().StorageType(),             \
+          lhs.quantized_per_tensor_element_type().ExpressedType(),           \
+          COMPARISON_OP, lhs, rhs, output)                                   \
+    }                                                                        \
+    break;                                                                   \
+  }
+
+  switch (op.attributes.comparison_direction) {
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kEq, std::equal_to<void>());
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kNe, std::not_equal_to<void>());
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kGe, std::greater_equal<void>());
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kGt, std::greater<void>());
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kLe, std::less_equal<void>());
+    SHLO_REF_COMPARISON_DIRECTION_CASE(kLt, std::less<void>());
+  }
+
+  return absl::FailedPreconditionError(
+      "stablehlo.compare: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/compare.h

@@ -0,0 +1,46 @@
+/* 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_OPS_COMPARE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COMPARE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CompareOp {
+  enum class ComparisonDirection { kEq, kNe, kGe, kGt, kLe, kLt };
+  // `compare_type` is considered for deprecation. We won't implement it until
+  // the deprecation is lifted.
+  //
+  // https://github.com/openxla/stablehlo/issues/584
+
+  struct Attributes {
+    ComparisonDirection comparison_direction;
+  };
+
+  Attributes attributes;
+};
+
+CompareOp Create(CompareOp::Attributes attributes);
+absl::Status Prepare(CompareOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output);
+absl::Status Evaluate(CompareOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COMPARE_H_

tensorflow/lite/experimental/shlo/ops/compare_test.cc

@@ -0,0 +1,257 @@
+/* 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/ops/compare.h"
+
+#include <string>
+#include <tuple>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/random/distributions.h"
+#include "absl/random/random.h"
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::FloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<CompareOp> {
+  static std::string Get() { return "Compare"; }
+};
+
+struct Compare {
+  template <class T>
+  constexpr bool operator()(const T a, const T b) const {
+    switch (comparison_direction) {
+      case CompareOp::ComparisonDirection::kEq:
+        return a == b;
+      case CompareOp::ComparisonDirection::kNe:
+        return a != b;
+      case CompareOp::ComparisonDirection::kGe:
+        return a >= b;
+      case CompareOp::ComparisonDirection::kGt:
+        return a > b;
+      case CompareOp::ComparisonDirection::kLe:
+        return a <= b;
+      case CompareOp::ComparisonDirection::kLt:
+        return a < b;
+    }
+    return false;
+  }
+
+  CompareOp::ComparisonDirection comparison_direction;
+};
+
+const char* ToString(CompareOp::ComparisonDirection comparison_direction) {
+  switch (comparison_direction) {
+    case CompareOp::ComparisonDirection::kEq:
+      return "eq";
+    case CompareOp::ComparisonDirection::kNe:
+      return "ne";
+    case CompareOp::ComparisonDirection::kGe:
+      return "ge";
+    case CompareOp::ComparisonDirection::kGt:
+      return "gt";
+    case CompareOp::ComparisonDirection::kLe:
+      return "le";
+    case CompareOp::ComparisonDirection::kLt:
+      return "lt";
+  }
+}
+
+template <>
+struct SupportedOpAttributes<CompareOp> {
+  static CompareOp::Attributes Get() {
+    return {.comparison_direction = CompareOp::ComparisonDirection::kEq};
+  }
+};
+
+template <>
+struct SupportedOpOutputDataType<CompareOp> {
+  static constexpr DataType kStorageType = DataType::kI1;
+};
+
+namespace {
+INSTANTIATE_TYPED_TEST_SUITE_P(Compare, BinaryElementwiseOpShapePropagationTest,
+                               CompareOp, TestParamNames);
+
+// Tests that the baseline element type of the input and output tensors is the
+// same.
+//
+// Compare has input/output constraints that are different from the other binary
+// element wise ops. Thus the specific test suite.
+GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(
+    BinaryElementwiseSameBaselineElementTypeConstraintTest);
+
+template <class Op, class SupportedTypes>
+using CompareBaselineConstraintTypesCrossProduct =
+    MapTypes<OpTupleFactory<Op>::template WithOp,
+             FilterTypes<NegatePred<SameTypes>::template Predicate,
+                         CrossProductTypes<SupportedTypes, SupportedTypes>>>;
+
+using CompareBaselineContraintTypes =
+    CompareBaselineConstraintTypesCrossProduct<
+        CompareOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes,
+                               BaselineConstraintFloatTypes,
+                               BaselineConstraintQuantizedPerTensorTypes>>;
+
+template <class T>
+class CompareSameBaselineElementTypeConstraintTest : public ::testing::Test {};
+
+TYPED_TEST_SUITE(CompareSameBaselineElementTypeConstraintTest,
+                 CompareBaselineContraintTypes, TestParamNames);
+
+TYPED_TEST(CompareSameBaselineElementTypeConstraintTest,
+           DifferentInputOutputStorageTypesRaiseAnError) {
+  using Op = std::tuple_element_t<0, TypeParam>;
+  using LhsTypeDesc = std::tuple_element_t<1, TypeParam>;
+  using RhsTypeDesc = std::tuple_element_t<2, TypeParam>;
+  const Shape shape({2, 3, 4});
+  Tensor lhs_tensor{.type = TensorTypeFor(LhsTypeDesc{}, shape),
+                    .data = nullptr};
+  Tensor rhs_tensor{.type = TensorTypeFor(RhsTypeDesc{}, shape),
+                    .data = nullptr};
+  Tensor output_tensor{.type = TensorTypeFor(TestParam<DataType::kI1>{}, shape),
+                       .data = nullptr};
+  auto op = Create(typename Op::Attributes{});
+  const absl::Status status =
+      Prepare(op, lhs_tensor, rhs_tensor, output_tensor);
+  EXPECT_THAT(status, shlo_ref::testing::StatusIs(
+                          absl::StatusCode::kFailedPrecondition));
+  EXPECT_THAT(
+      status.message(),
+      ::testing::ContainsRegex(
+          "stablehlo.[_a-z]+: baseline type constraint is not satisfied"));
+}
+
+using UnsupportedTypes =
+    WithOpTypes<CompareOp, ConcatTypes<PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Compare, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, ArithmeticTestTypes>;
+
+template <class T>
+struct CompareTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(CompareTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(CompareTest, SupportedTestTypesTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  absl::BitGen bit_gen;
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-50, /*max=*/50);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/1, /*max=*/5);
+  Vector<StorageT> output_data(shape.NumElements());
+  Tensor lhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = DataType::kI1},
+      .data = output_data.data()};
+
+  const CompareOp::ComparisonDirection comparison_direction =
+      static_cast<CompareOp::ComparisonDirection>(absl::Uniform(bit_gen, 0, 6));
+
+  Compare compare_ref{comparison_direction};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(lhs_data, rhs_data, expected_data.begin(), compare_ref);
+
+  auto op = Create(
+      CompareOp::Attributes{.comparison_direction = comparison_direction});
+  ASSERT_OK(Prepare(op, lhs_tensor, rhs_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, lhs_tensor, rhs_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedCompareTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedCompareTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedCompareTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  absl::BitGen bit_gen;
+  const Shape shape({2, 2, 2});
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(2);
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-50, /*max=*/50);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/zero_point + 1,
+                                        /*max=*/zero_point + 5);
+  Vector<StorageType<DataType::kI1>> output_data(shape.NumElements());
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor lhs_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = DataType::kI1},
+      .data = output_data.data()};
+
+  const CompareOp::ComparisonDirection comparison_direction =
+      CompareOp::ComparisonDirection::kEq;
+  // static_cast<CompareOp::ComparisonDirection>(absl::Uniform(bit_gen,
+  // 0, 6));
+
+  Compare compare_ref{comparison_direction};
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      lhs_data, rhs_data, expected_data.begin(),
+      [zero_point, scale, compare_ref](auto lhs, auto rhs) {
+        const ExpressedT dequantized_lhs = Dequantize(lhs, zero_point, scale);
+        const ExpressedT dequantized_rhs = Dequantize(rhs, zero_point, scale);
+        return compare_ref(dequantized_lhs, dequantized_rhs);
+      });
+
+  auto op = Create(
+      CompareOp::Attributes{.comparison_direction = comparison_direction});
+  ASSERT_OK(Prepare(op, lhs_tensor, rhs_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, lhs_tensor, rhs_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data))
+      << "lhs " << ::testing::PrintToString(lhs_data) << "\n"
+      << "rhs " << ::testing::PrintToString(rhs_data) << "\n"
+      << "dir " << ToString(comparison_direction) << "\n";
+}
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/cosine.cc

@@ -0,0 +1,73 @@
+/* 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/ops/cosine.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Cosine {
+  template <class T>
+  T operator()(T v) const {
+    return std::cos(v);
+  }
+};
+
+template <>
+F16 Cosine::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Cosine::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+CosineOp Create(CosineOp::Attributes) { return {}; }
+
+absl::Status Prepare(CosineOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("cosine"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("cosine"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(CosineOp& op, const Tensor& input, Tensor& output) {
+  Cosine cosine;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), cosine,
+        input, output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   cosine, input, output);
+  }
+  return absl::FailedPreconditionError("Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/cosine.h

@@ -0,0 +1,34 @@
+/* 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_OPS_COSINE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COSINE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CosineOp {
+  struct Attributes {};
+};
+
+CosineOp Create(CosineOp::Attributes);
+absl::Status Prepare(CosineOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(CosineOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COSINE_H_

tensorflow/lite/experimental/shlo/ops/cosine_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 "tensorflow/lite/experimental/shlo/ops/cosine.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<CosineOp> {
+  static std::string Get() { return "Cosine"; }
+};
+
+namespace {
+
+struct Cosine {
+  template <class T>
+  T operator()(T v) const {
+    return std::cos(v);
+  }
+} cosine_ref;
+
+template <>
+F16 Cosine::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Cosine::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Cosine, UnaryElementwiseOpShapePropagationTest,
+                               CosineOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Cosine, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<CosineOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<CosineOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                      PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Cosine, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct CosineTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(CosineTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(CosineTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), cosine_ref);
+
+  auto op = Create(CosineOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedCosineTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedCosineTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedCosineTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = cosine_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(CosineOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/count_leading_zeros.cc

@@ -0,0 +1,65 @@
+/* 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/ops/count_leading_zeros.h"
+
+#include <cstdint>
+#include <type_traits>
+
+#include "absl/numeric/bits.h"
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/i4.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CountLeadingZeros {
+  template <class T>
+  T operator()(T v) const {
+    if constexpr (std::is_same_v<I4, T>) {
+      return I4(absl::countl_zero(static_cast<uint8_t>(v << 4 | 0xf)));
+    } else {
+      return absl::countl_zero(static_cast<std::make_unsigned_t<T>>(v));
+    }
+  }
+};
+
+CountLeadingZerosOp Create(CountLeadingZerosOp::Attributes) { return {}; }
+
+absl::Status Prepare(CountLeadingZerosOp& op, const Tensor& input,
+                     Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("count_leading_zeros"), input, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("count_leading_zeros"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(CountLeadingZerosOp& op, const Tensor& input,
+                      Tensor& output) {
+  CountLeadingZeros count_leading_zeros;
+  if (IsIntTensor(input)) {
+    DISPATCH_INT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                 count_leading_zeros, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.count_leading_zeros: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/count_leading_zeros.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_OPS_COUNT_LEADING_ZEROS_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COUNT_LEADING_ZEROS_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct CountLeadingZerosOp {
+  struct Attributes {};
+};
+
+CountLeadingZerosOp Create(CountLeadingZerosOp::Attributes);
+absl::Status Prepare(CountLeadingZerosOp& op, const Tensor& input,
+                     Tensor& output);
+absl::Status Evaluate(CountLeadingZerosOp& op, const Tensor& input,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_COUNT_LEADING_ZEROS_H_

tensorflow/lite/experimental/shlo/ops/count_leading_zeros_test.cc

@@ -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.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/ops/count_leading_zeros.h"
+
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <type_traits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/numeric/bits.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/i4.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<CountLeadingZerosOp> {
+  static std::string Get() { return "CountLeadingZeros"; }
+};
+
+template <>
+struct SupportedOpDataType<CountLeadingZerosOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+struct CountLeadingZeros {
+  template <class T>
+  T operator()(T v) const {
+    if constexpr (std::is_same_v<I4, T>) {
+      return I4(absl::countl_zero(static_cast<uint8_t>(v << 4 | 0xf)));
+    } else {
+      return absl::countl_zero(static_cast<std::make_unsigned_t<T>>(v));
+    }
+  }
+} count_leading_zeros_ref;
+
+template <class T>
+struct CountLeadingZerosFunctorTest : ::testing::Test {};
+
+using CountLeadingZerosTypes = ::testing::Types<int32_t, int16_t, int8_t, I4>;
+
+TYPED_TEST_SUITE(CountLeadingZerosFunctorTest, CountLeadingZerosTypes);
+
+TYPED_TEST(CountLeadingZerosFunctorTest, GivesCorrectResults) {
+  int64_t bit_count = 8 * sizeof(TypeParam);
+  if constexpr (std::is_same_v<I4, TypeParam>) {
+    bit_count = 4;
+  }
+  EXPECT_EQ(count_leading_zeros_ref(std::numeric_limits<TypeParam>::lowest()),
+            0);
+  EXPECT_EQ(count_leading_zeros_ref(static_cast<TypeParam>(-1)), 0);
+  EXPECT_EQ(count_leading_zeros_ref(static_cast<TypeParam>(0)), bit_count);
+  EXPECT_EQ(count_leading_zeros_ref(static_cast<TypeParam>(1)), bit_count - 1);
+  EXPECT_EQ(count_leading_zeros_ref(static_cast<TypeParam>(2)), bit_count - 2);
+  EXPECT_EQ(count_leading_zeros_ref(std::numeric_limits<TypeParam>::max()), 1);
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(CountLeadingZeros,
+                               UnaryElementwiseOpShapePropagationTest,
+                               CountLeadingZerosOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    CountLeadingZeros, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    BaselineMismatchSignedIntegerTypes<CountLeadingZerosOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<CountLeadingZerosOp, ConcatTypes<BoolTestType, FloatTestTypes,
+                                                 PerTensorQuantizedTestTypes,
+                                                 PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(CountLeadingZeros,
+                               UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct CountLeadingZerosTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(CountLeadingZerosTest, IntTestTypes, TestParamNames);
+
+TYPED_TEST(CountLeadingZerosTest, IntTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = IotaBuffer<TypeParam::kStorage>(shape, -12);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), count_leading_zeros_ref);
+
+  auto op = Create(CountLeadingZerosOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/divide.cc

@@ -0,0 +1,62 @@
+/* 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/ops/divide.h"
+
+#include <functional>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Divide : std::divides<void> {};
+
+DivideOp Create(DivideOp::Attributes) { return {}; }
+
+absl::Status Prepare(DivideOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(lhs.shape(), rhs.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(CheckCtx("divide"), lhs,
+                                               IsIntTensor, IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("divide"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("divide"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(DivideOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  Divide divide;
+  if (IsIntTensor(lhs) || IsFloatTensor(lhs)) {
+    // Note: all the arithmetic types share the same implementation.
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       lhs.tensor_element_type(), divide, lhs, rhs, output);
+  } else if (IsQuantizedPerTensorTensor(lhs)) {
+    DISPATCH_QUANTIZED(detail::DequantizeOpQuantizePerTensor,
+                       lhs.quantized_per_tensor_element_type().StorageType(),
+                       lhs.quantized_per_tensor_element_type().ExpressedType(),
+                       divide, lhs, rhs, output)
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.divide: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/divide.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_OPS_DIVIDE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_DIVIDE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct DivideOp {
+  struct Attributes {};
+};
+
+DivideOp Create(DivideOp::Attributes);
+absl::Status Prepare(DivideOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output);
+absl::Status Evaluate(DivideOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_DIVIDE_H_

tensorflow/lite/experimental/shlo/ops/divide_test.cc

@@ -0,0 +1,149 @@
+/* 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/ops/divide.h"
+
+#include <functional>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/ops/binary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::FloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<DivideOp> {
+  static std::string Get() { return "Divide"; }
+};
+
+struct Divide : std::divides<void> {};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Divide, BinaryElementwiseOpShapePropagationTest,
+                               DivideOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    DivideOp,
+    ConcatTypes<BaselineConstraintIntTypes, BaselineConstraintFloatTypes,
+                BaselineConstraintQuantizedPerTensorTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Divide, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<DivideOp, ConcatTypes<BoolTestType, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Divide, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using ArithmeticTypes = ConcatTypes<ArithmeticTestTypes>;
+
+template <class T>
+struct DivideTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(DivideTest, ArithmeticTypes, TestParamNames);
+
+TYPED_TEST(DivideTest, ArithmeticTestTypesTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-50, /*max=*/50);
+  Vector<StorageT> rhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/1, /*max=*/5);
+  Vector<StorageT> output_data(shape.NumElements());
+  Tensor lhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(lhs_data, rhs_data, expected_data.begin(), Divide());
+
+  auto op = Create(DivideOp::Attributes{});
+  ASSERT_OK(Prepare(op, lhs_tensor, rhs_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, lhs_tensor, rhs_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedDivideTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedDivideTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedDivideTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(2);
+  Vector<StorageT> lhs_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/-50, /*max=*/50);
+  Vector<StorageT> rhs_data = RandomBuffer<TypeParam::kStorage>(
+      shape, /*min=*/zero_point + 1, /*max=*/zero_point + 5);
+  Vector<StorageT> output_data(shape.NumElements());
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor lhs_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = lhs_data.data()};
+  Tensor rhs_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = rhs_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      lhs_data, rhs_data, expected_data.begin(),
+      [zero_point, scale](auto lhs, auto rhs) {
+        const ExpressedT dequantized_lhs = Dequantize(lhs, zero_point, scale);
+        const ExpressedT dequantized_rhs = Dequantize(rhs, zero_point, scale);
+        const ExpressedT dequantized_res =
+            Divide()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(DivideOp::Attributes{});
+  ASSERT_OK(Prepare(op, lhs_tensor, rhs_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, lhs_tensor, rhs_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(FloatEq(), expected_data));
+}
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/exponential.cc

@@ -0,0 +1,74 @@
+/* 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/ops/exponential.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Exponential {
+  template <class T>
+  T operator()(T v) const {
+    return std::exp(v);
+  }
+};
+
+template <>
+F16 Exponential::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Exponential::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+ExponentialOp Create(ExponentialOp::Attributes) { return {}; }
+
+absl::Status Prepare(ExponentialOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("cosine"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("cosine"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(ExponentialOp& op, const Tensor& input, Tensor& output) {
+  Exponential exponential;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), exponential,
+        input, output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   exponential, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.tanh: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/exponential.h

@@ -0,0 +1,34 @@
+/* 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_OPS_EXPONENTIAL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_EXPONENTIAL_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct ExponentialOp {
+  struct Attributes {};
+};
+
+ExponentialOp Create(ExponentialOp::Attributes);
+absl::Status Prepare(ExponentialOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(ExponentialOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_EXPONENTIAL_H_

tensorflow/lite/experimental/shlo/ops/exponential_minus_one.cc

@@ -0,0 +1,77 @@
+/* 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/ops/exponential_minus_one.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct ExponentialMinusOne {
+  template <class T>
+  T operator()(T v) const {
+    return std::expm1(v);
+  }
+};
+
+template <>
+F16 ExponentialMinusOne::operator()(F16 v) const {
+  return F16(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 ExponentialMinusOne::operator()(BF16 v) const {
+  return BF16(operator()(static_cast<float>(v)));
+}
+
+ExponentialMinusOneOp Create(ExponentialMinusOneOp::Attributes) { return {}; }
+
+absl::Status Prepare(ExponentialMinusOneOp& op, const Tensor& input,
+                     Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("exponential_minus_one"), input,
+                          IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("exponential_minus_one"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(ExponentialMinusOneOp& op, const Tensor& input,
+                      Tensor& output) {
+  ExponentialMinusOne exponential_minus_one;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(),
+        exponential_minus_one, input, output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   exponential_minus_one, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.exponential_minus_one: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/exponential_minus_one.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_OPS_EXPONENTIAL_MINUS_ONE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_EXPONENTIAL_MINUS_ONE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct ExponentialMinusOneOp {
+  struct Attributes {};
+};
+
+ExponentialMinusOneOp Create(ExponentialMinusOneOp::Attributes);
+absl::Status Prepare(ExponentialMinusOneOp& op, const Tensor& input,
+                     Tensor& output);
+absl::Status Evaluate(ExponentialMinusOneOp& op, const Tensor& input,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_EXPONENTIAL_MINUS_ONE_H_

tensorflow/lite/experimental/shlo/ops/exponential_minus_one_test.cc

@@ -0,0 +1,152 @@
+/* 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/ops/exponential_minus_one.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<ExponentialMinusOneOp> {
+  static std::string Get() { return "ExponentialMinusOne"; }
+};
+
+namespace {
+
+struct ExponentialMinusOne {
+  template <class T>
+  T operator()(T v) const {
+    return std::expm1(v);
+  }
+} exponential_minus_one_ref;
+
+template <>
+F16 ExponentialMinusOne::operator()(F16 v) const {
+  return F16(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 ExponentialMinusOne::operator()(BF16 v) const {
+  return BF16(operator()(static_cast<float>(v)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(ExponentialMinusOne,
+                               UnaryElementwiseOpShapePropagationTest,
+                               ExponentialMinusOneOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    ExponentialMinusOne, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<ExponentialMinusOneOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<ExponentialMinusOneOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                                   PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(ExponentialMinusOneOp,
+                               UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct ExponentialMinusOneTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(ExponentialMinusOneTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(ExponentialMinusOneTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(),
+                    exponential_minus_one_ref);
+
+  auto op = Create(ExponentialMinusOneOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedExponentialMinusOneTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedExponentialMinusOneTest, QuantizedTestTypes,
+                 TestParamNames);
+
+TYPED_TEST(QuantizedExponentialMinusOneTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res =
+            exponential_minus_one_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(ExponentialMinusOneOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/exponential_test.cc

@@ -0,0 +1,148 @@
+/* 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/ops/exponential.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<ExponentialOp> {
+  static std::string Get() { return "Exponential"; }
+};
+
+namespace {
+
+struct Exponential {
+  template <class T>
+  T operator()(T v) const {
+    return std::exp(v);
+  }
+} exponential_ref;
+
+template <>
+F16 Exponential::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Exponential::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Exponential,
+                               UnaryElementwiseOpShapePropagationTest,
+                               ExponentialOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Exponential, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<ExponentialOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<ExponentialOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                           PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Exponential, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct ExponentialTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(ExponentialTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(ExponentialTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), exponential_ref);
+
+  auto op = Create(ExponentialOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedExponentialTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedExponentialTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedExponentialTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = exponential_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(ExponentialOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/floor.cc

@@ -0,0 +1,74 @@
+/* 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/ops/floor.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Floor {
+  template <class T>
+  T operator()(T v) const {
+    return std::floor(v);
+  }
+};
+
+template <>
+F16 Floor::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Floor::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+FloorOp Create(FloorOp::Attributes) { return {}; }
+
+absl::Status Prepare(FloorOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("floor"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("floor"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(FloorOp& op, const Tensor& input, Tensor& output) {
+  Floor floor;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), floor, input,
+        output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   floor, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.floor: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/floor.h

@@ -0,0 +1,34 @@
+/* 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_OPS_FLOOR_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_FLOOR_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct FloorOp {
+  struct Attributes {};
+};
+
+FloorOp Create(FloorOp::Attributes);
+absl::Status Prepare(FloorOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(FloorOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_FLOOR_H_

tensorflow/lite/experimental/shlo/ops/floor_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 "tensorflow/lite/experimental/shlo/ops/floor.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<FloorOp> {
+  static std::string Get() { return "Floor"; }
+};
+
+namespace {
+
+struct Floor {
+  template <class T>
+  T operator()(T v) const {
+    return std::floor(v);
+  }
+} floor_ref;
+
+template <>
+F16 Floor::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Floor::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Floor, UnaryElementwiseOpShapePropagationTest,
+                               FloorOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Floor, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<FloorOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<FloorOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                     PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Floor, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct FloorTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(FloorTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(FloorTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), floor_ref);
+
+  auto op = Create(FloorOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedFloorTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedFloorTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedFloorTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = floor_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(FloorOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/is_finite.cc

@@ -0,0 +1,97 @@
+/* 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/ops/is_finite.h"
+
+#include <algorithm>
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+namespace {
+
+absl::Status CheckParameters(const Tensor& operand, const Tensor& result) {
+  if (operand.shape() != result.shape()) {
+    return absl::InvalidArgumentError(
+        "operand and result must have the same shape");
+  }
+  if (!operand.IsQuantized() && !IsFloat(operand.tensor_element_type())) {
+    return absl::InvalidArgumentError(
+        "operand must be floating-point type or per-tensor quantized");
+  }
+  if (operand.IsPerAxisQuantized()) {
+    return absl::InvalidArgumentError("operand cannot be per-axis quantized");
+  }
+  if (result.IsQuantized()) {
+    return absl::InvalidArgumentError("result cannot be quantized");
+  }
+  if (!IsBool(result.tensor_element_type())) {
+    return absl::InvalidArgumentError("result must be an I1 tensor");
+  }
+  if (operand.NumElements() != result.NumElements()) {
+    return absl::InvalidArgumentError(
+        "operand and result must have the same size");
+  }
+  return absl::OkStatus();
+}
+
+template <DataType data_type>
+absl::Status EvaluateImpl(const Tensor& operand, bool* output) {
+  const auto* in = operand.GetDataAs<data_type>();
+  const auto num_elements = operand.NumElements();
+  for (DimensionSize i = 0; i < num_elements; ++i) {
+    output[i] = std::isfinite(static_cast<float>(in[i]));
+  }
+  return absl::OkStatus();
+}
+
+absl::Status EvaluateImpl(const Tensor& operand, Tensor& result) {
+  bool* output = result.GetDataAs<DataType::kI1>();
+  if (!operand.IsQuantized()) {
+    DISPATCH_FLOAT(EvaluateImpl, operand.tensor_element_type(), operand,
+                   output);
+  } else {  // IsQuantized(operand)
+    // For quantized types, the result is always true.
+    const auto num_elements = result.NumElements();
+    std::fill(output, output + num_elements, true);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+IsFiniteOp Create(const IsFiniteOp::Attributes& attributes) {
+  return IsFiniteOp();
+}
+
+absl::Status Prepare(IsFiniteOp& op, const Tensor& operand, Tensor& result) {
+  return CheckParameters(operand, result);
+}
+
+absl::Status Evaluate(IsFiniteOp& op, const Tensor& operand, Tensor& result) {
+  if (!operand.data) {
+    return absl::InvalidArgumentError("No operand.data");
+  }
+  if (!result.data) {
+    return absl::InvalidArgumentError("No result.data");
+  }
+  return EvaluateImpl(operand, result);
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/is_finite.h

@@ -0,0 +1,34 @@
+/* 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_OPS_IS_FINITE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_IS_FINITE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct IsFiniteOp {
+  struct Attributes {};
+};
+
+IsFiniteOp Create(const IsFiniteOp::Attributes& attributes);
+absl::Status Prepare(IsFiniteOp& op, const Tensor& operand, Tensor& result);
+absl::Status Evaluate(IsFiniteOp& op, const Tensor& operand, Tensor& result);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_IS_FINITE_H_

tensorflow/lite/experimental/shlo/ops/is_finite_bench.cc

@@ -0,0 +1,87 @@
+/* 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 <vector>
+
+#include "absl/log/absl_check.h"
+#include "benchmark/benchmark.h"  // from @com_google_benchmark
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/ops/benchmark_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/is_finite.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+#include "tensorflow/lite/experimental/shlo/tensor_with_data.h"
+
+namespace shlo_ref {
+namespace {
+
+void BM_IsFinite(benchmark::State& state, DimensionSize num_elements,
+                 const Tensor& operand) {
+  IsFiniteOp op = Create(IsFiniteOp::Attributes{});
+
+  Tensor result = Tensor{.type = TensorType{.shape = Shape{{num_elements}},
+                                            .element_type = DataType::kI1}};
+  ABSL_CHECK_OK(Prepare(op, operand, result));
+
+  std::vector<std::byte> result_values(result.SizeInBytes());
+  result.data = result_values.data();
+
+  for (auto _ : state) {
+    ABSL_CHECK_OK(Evaluate(op, operand, result));
+  }
+}
+
+template <DataType data_type>
+void BM_IsFinite(benchmark::State& state) {
+  const DimensionSize num_elements = state.range(0);
+
+  auto operand_values = GenerateRandomVector<data_type>(num_elements);
+  auto operand =
+      TensorWithData::Create<data_type>(Shape{{num_elements}}, operand_values);
+
+  BM_IsFinite(state, num_elements, operand.tensor());
+}
+
+template <DataType storage_type, DataType expressed_type>
+void BM_IsFiniteQuantized(benchmark::State& state) {
+  const DimensionSize num_elements = state.range(0);
+
+  auto operand_values = GenerateRandomVector<expressed_type>(num_elements);
+  auto operand = TensorWithData::Create<storage_type, expressed_type>(
+      Shape{{num_elements}}, operand_values, 0.1, 0);
+
+  BM_IsFinite(state, num_elements, operand.tensor());
+}
+
+BENCHMARK(BM_IsFinite<DataType::kBF16>)
+    ->RangeMultiplier(2)
+    ->Range(KiB(8), KiB(64));
+BENCHMARK(BM_IsFinite<DataType::kF16>)
+    ->RangeMultiplier(2)
+    ->Range(KiB(8), KiB(64));
+BENCHMARK(BM_IsFinite<DataType::kF32>)
+    ->RangeMultiplier(2)
+    ->Range(KiB(8), KiB(64));
+
+// IsFinite will be the same regardless of quantization parameters, so only
+// benchmark one combination.
+BENCHMARK(BM_IsFiniteQuantized<DataType::kSI16, DataType::kF32>)
+    ->RangeMultiplier(2)
+    ->Range(KiB(8), KiB(64));
+
+}  // namespace
+}  // namespace shlo_ref
+
+BENCHMARK_MAIN();

tensorflow/lite/experimental/shlo/ops/is_finite_test.cc

@@ -0,0 +1,90 @@
+/* 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/ops/is_finite.h"
+
+#include <cmath>
+#include <cstddef>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+#include "tensorflow/lite/experimental/shlo/tensor_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor_with_data.h"
+
+namespace shlo_ref {
+namespace {
+
+using ::shlo_ref::testing::TensorEq;
+
+struct Params {
+  TensorWithData operand;
+  TensorWithData expected;
+};
+
+class IsFiniteTest : public ::testing::TestWithParam<Params> {};
+
+TEST_P(IsFiniteTest, IsFinite) {
+  const auto& params = GetParam();
+
+  IsFiniteOp op = Create(IsFiniteOp::Attributes{});
+  Tensor result{.type = params.expected.tensor().type};
+
+  ASSERT_OK(Prepare(op, params.operand.tensor(), result));
+
+  std::vector<std::byte> result_data(result.SizeInBytes());
+  result.data = result_data.data();
+
+  EXPECT_OK(Evaluate(op, params.operand.tensor(), result));
+  EXPECT_THAT(result, TensorEq(params.expected.tensor()));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Unquantized, IsFiniteTest,
+    ::testing::Values(
+        Params{TensorWithData::Create<DataType::kBF16>(
+                   Shape{{7}},
+                   {BF16{+NAN}, BF16{-NAN}, BF16{-INFINITY}, BF16{+INFINITY},
+                    BF16{-1.0f}, BF16{0.0f}, BF16{1.0f}}),
+               TensorWithData::Create<DataType::kI1>(
+                   Shape{{7}}, {false, false, false, false, true, true, true})},
+        Params{
+            TensorWithData::Create<DataType::kF16>(
+                Shape{{7}}, {F16{+NAN}, F16{-NAN}, F16{-INFINITY},
+                             F16{+INFINITY}, F16{-1.0f}, F16{0.0f}, F16{1.0f}}),
+            TensorWithData::Create<DataType::kI1>(
+                Shape{{7}}, {false, false, false, false, true, true, true})},
+        Params{
+            TensorWithData::Create<DataType::kF32>(
+                Shape{{7}},
+                {+NAN, -NAN, -INFINITY, +INFINITY, -1.0f, 0.0f, 1.0f}),
+            TensorWithData::Create<DataType::kI1>(
+                Shape{{7}}, {false, false, false, false, true, true, true})}));
+
+INSTANTIATE_TEST_SUITE_P(
+    Quantized, IsFiniteTest,
+    ::testing::Values(Params{
+        .operand = TensorWithData::Create<DataType::kSI16, DataType::kF32>(
+            Shape{{7}}, {0.0f, -1.0f, 0.0f, 1.0f, 1.0f, 0.0f, 1.0f}, 0.1f, 0),
+        .expected = TensorWithData::Create<DataType::kI1>(
+            Shape{{7}}, {true, true, true, true, true, true, true})}));
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/log.cc

@@ -0,0 +1,74 @@
+/* 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/ops/log.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Log {
+  template <class T>
+  T operator()(T v) const {
+    return std::log(v);
+  }
+};
+
+template <>
+F16 Log::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Log::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+LogOp Create(LogOp::Attributes) { return {}; }
+
+absl::Status Prepare(LogOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("log"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("log"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(LogOp& op, const Tensor& input, Tensor& output) {
+  Log log;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), log, input,
+        output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   log, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.log: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/log.h

@@ -0,0 +1,34 @@
+/* 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_OPS_LOG_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_LOG_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct LogOp {
+  struct Attributes {};
+};
+
+LogOp Create(LogOp::Attributes);
+absl::Status Prepare(LogOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(LogOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_LOG_H_

tensorflow/lite/experimental/shlo/ops/log_plus_one.cc

@@ -0,0 +1,75 @@
+/* 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/ops/log_plus_one.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct LogPlusOne {
+  template <class T>
+  T operator()(T v) const {
+    return std::log1p(v);
+  }
+};
+
+template <>
+F16 LogPlusOne::operator()(F16 v) const {
+  return F16(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 LogPlusOne::operator()(BF16 v) const {
+  return BF16(operator()(static_cast<float>(v)));
+}
+
+LogPlusOneOp Create(LogPlusOneOp::Attributes) { return {}; }
+
+absl::Status Prepare(LogPlusOneOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(CheckCtx("log_plus_one"), input,
+                                               IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("log_plus_one"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(LogPlusOneOp& op, const Tensor& input, Tensor& output) {
+  LogPlusOne log_plus_one;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), log_plus_one,
+        input, output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   log_plus_one, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.log_plus_one: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/log_plus_one.h

@@ -0,0 +1,34 @@
+/* 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_OPS_LOG_PLUS_ONE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_LOG_PLUS_ONE_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct LogPlusOneOp {
+  struct Attributes {};
+};
+
+LogPlusOneOp Create(LogPlusOneOp::Attributes);
+absl::Status Prepare(LogPlusOneOp& op, const Tensor& input, Tensor& output);
+absl::Status Evaluate(LogPlusOneOp& op, const Tensor& input, Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_LOG_PLUS_ONE_H_

tensorflow/lite/experimental/shlo/ops/log_plus_one_test.cc

@@ -0,0 +1,150 @@
+/* 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/ops/log_plus_one.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<LogPlusOneOp> {
+  static std::string Get() { return "LogPlusOne"; }
+};
+
+namespace {
+
+struct LogPlusOne {
+  template <class T>
+  T operator()(T v) const {
+    return std::log1p(v);
+  }
+} log_plus_one_ref;
+
+template <>
+F16 LogPlusOne::operator()(F16 v) const {
+  return F16(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 LogPlusOne::operator()(BF16 v) const {
+  return BF16(operator()(static_cast<float>(v)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(LogPlusOne,
+                               UnaryElementwiseOpShapePropagationTest,
+                               LogPlusOneOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    LogPlusOne, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<LogPlusOneOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<LogPlusOneOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                          PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(LogPlusOne, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct LogPlusOneTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(LogPlusOneTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(LogPlusOneTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(
+      shape, /*min=*/static_cast<StorageT>(-0.99));
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), log_plus_one_ref);
+
+  auto op = Create(LogPlusOneOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedLogPlusOneTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedLogPlusOneTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedLogPlusOneTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  Vector<StorageT> input_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/zero_point);
+  Vector<StorageT> output_data(shape.NumElements());
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = log_plus_one_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(LogPlusOneOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/log_test.cc

@@ -0,0 +1,148 @@
+/* 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/ops/log.h"
+
+#include <cmath>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/test_util.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise_test_util.h"
+#include "tensorflow/lite/experimental/shlo/quantize.h"
+#include "tensorflow/lite/experimental/shlo/quantized_tensor_element_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<LogOp> {
+  static std::string Get() { return "Log"; }
+};
+
+namespace {
+
+struct Log {
+  template <class T>
+  T operator()(T v) const {
+    return std::log(v);
+  }
+} log_ref;
+
+template <>
+F16 Log::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Log::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Log, UnaryElementwiseOpShapePropagationTest,
+                               LogOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Log, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<LogOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    LogOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Log, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct LogTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(LogTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(LogTest, FloatTensorsWork) {
+  using StorageT = typename TypeParam::StorageT;
+
+  const Shape shape({2, 3, 4});
+  Vector<StorageT> input_data = RandomBuffer<TypeParam::kStorage>(
+      shape, /*min=*/static_cast<StorageT>(0.1));
+  Vector<StorageT> output_data(shape.NumElements());
+
+  Tensor input_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = TensorType{.shape = shape, .element_type = TypeParam::kStorage},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), log_ref);
+
+  auto op = Create(LogOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, Pointwise(NanSensitiveFloatEq(), expected_data));
+}
+
+template <class T>
+struct QuantizedLogTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedLogTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedLogTest, PerTensorWorks) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({2, 3, 4});
+  const ExpressedT scale = static_cast<ExpressedT>(1.5);
+  const StorageT zero_point = static_cast<StorageT>(5);
+  Vector<StorageT> input_data =
+      RandomBuffer<TypeParam::kStorage>(shape, /*min=*/zero_point + 1);
+  Vector<StorageT> output_data(shape.NumElements());
+  const QuantizedElementTypePerTensor tensor_type =
+      QuantizedElementTypePerTensor(TypeParam::kStorage, zero_point,
+                                    TypeParam::kExpressed, scale);
+  Tensor input_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerTensorTensorType{.shape = shape,
+                                           .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(), [zero_point, scale](auto v) {
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = log_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(LogOp::Attributes{});
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+}  // namespace
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/logistic.cc

@@ -0,0 +1,75 @@
+/* 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/ops/logistic.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/f16.h"
+#include "tensorflow/lite/experimental/shlo/ops/unary_elementwise.h"
+#include "tensorflow/lite/experimental/shlo/ops/util.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct Logistic {
+  template <class T>
+  T operator()(T v) const {
+    constexpr T one = static_cast<T>(1);
+    return one / (one + std::exp(-v));
+  }
+};
+
+template <>
+F16 Logistic::operator()(F16 v) const {
+  return F16(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 Logistic::operator()(BF16 v) const {
+  return BF16(operator()(static_cast<float>(v)));
+}
+
+LogisticOp Create(LogisticOp::Attributes) { return {}; }
+
+absl::Status Prepare(LogisticOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("logistic"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("logistic"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(LogisticOp& op, const Tensor& input, Tensor& output) {
+  Logistic logistic;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), logistic,
+        input, output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   logistic, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.logistic: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref