chore: import upstream snapshot with attribution

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

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

tensorflow/lite/experimental/shlo/ops/logistic_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/logistic.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<LogisticOp> {
+  static std::string Get() { return "Logistic"; }
+};
+
+namespace {
+
+struct Logistic {
+  template <class T>
+  T operator()(T v) const {
+    constexpr T one = static_cast<T>(1);
+    return one / (one + std::exp(-v));
+  }
+} logistic_ref;
+
+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)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Logistic, UnaryElementwiseOpShapePropagationTest,
+                               LogisticOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Logistic, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<LogisticOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<LogisticOp, ConcatTypes<BoolTestType, IntTestTypes,
+                                        PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Logistic, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct LogisticTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(LogisticTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(LogisticTest, 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(), logistic_ref);
+
+  auto op = Create(LogisticOp::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 QuantizedLogisticTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedLogisticTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedLogisticTest, 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 = logistic_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(LogisticOp::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/maximum.cc

@@ -0,0 +1,66 @@
+/* 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/maximum.h"
+
+#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 Maximum {
+  template <class T>
+  constexpr auto operator()(const T a, const T b) {
+    return a > b ? a : b;
+  }
+};
+
+MaximumOp Create(MaximumOp::Attributes) { return {}; }
+
+absl::Status Prepare(MaximumOp& 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("maximum"), lhs, IsBoolTensor, IsIntTensor,
+                          IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("maximum"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("maximum"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(MaximumOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  Maximum maximum;
+  if (IsBoolTensor(lhs) || IsIntTensor(lhs) || IsFloatTensor(lhs)) {
+    // Note: all the arithmetic types share the same implementation.
+    DISPATCH_BOOL_INT_FLOAT(detail::EvaluateNoQuantization,
+                            lhs.tensor_element_type(), maximum, 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(),
+                       maximum, lhs, rhs, output)
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.maximum: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

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

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

@@ -0,0 +1,153 @@
+/* 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/maximum.h"
+
+#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<MaximumOp> {
+  static std::string Get() { return "Maximum"; }
+};
+
+struct Maximum {
+  template <class T>
+  constexpr auto operator()(const T a, const T b) {
+    return a > b ? a : b;
+  }
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Maximum, BinaryElementwiseOpShapePropagationTest,
+                               MaximumOp, TestParamNames);
+
+using MaximumBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    MaximumOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes,
+                           BaselineConstraintFloatTypes,
+                           BaselineConstraintQuantizedPerTensorTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Maximum, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MaximumBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<MaximumOp, ConcatTypes<PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Maximum, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, ArithmeticTestTypes>;
+
+template <class T>
+struct MaximumTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(MaximumTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(MaximumTest, 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(), Maximum());
+
+  auto op = Create(MaximumOp::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 QuantizedMaximumTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedMaximumTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedMaximumTest, 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 =
+            Maximum()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(MaximumOp::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/minimum.cc

@@ -0,0 +1,66 @@
+/* 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/minimum.h"
+
+#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 Minimum {
+  template <class T>
+  constexpr auto operator()(const T a, const T b) {
+    return a < b ? a : b;
+  }
+};
+
+MinimumOp Create(MinimumOp::Attributes) { return {}; }
+
+absl::Status Prepare(MinimumOp& 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("minimum"), lhs, IsBoolTensor, IsIntTensor,
+                          IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("minimum"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("minimum"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(MinimumOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  Minimum minimum;
+  if (IsBoolTensor(lhs) || IsIntTensor(lhs) || IsFloatTensor(lhs)) {
+    // Note: all the arithmetic types share the same implementation.
+    DISPATCH_BOOL_INT_FLOAT(detail::EvaluateNoQuantization,
+                            lhs.tensor_element_type(), minimum, 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(),
+                       minimum, lhs, rhs, output)
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.minimum: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

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

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

@@ -0,0 +1,153 @@
+/* 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/minimum.h"
+
+#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<MinimumOp> {
+  static std::string Get() { return "Minimum"; }
+};
+
+struct Minimum {
+  template <class T>
+  constexpr auto operator()(const T a, const T b) {
+    return a < b ? a : b;
+  }
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Minimum, BinaryElementwiseOpShapePropagationTest,
+                               MinimumOp, TestParamNames);
+
+using MinimumBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    MinimumOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes,
+                           BaselineConstraintFloatTypes,
+                           BaselineConstraintQuantizedPerTensorTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Minimum, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MinimumBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<MinimumOp, ConcatTypes<PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Minimum, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, ArithmeticTestTypes>;
+
+template <class T>
+struct MinimumTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(MinimumTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(MinimumTest, 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(), Minimum());
+
+  auto op = Create(MinimumOp::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 QuantizedMinimumTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedMinimumTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedMinimumTest, 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 =
+            Minimum()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(MinimumOp::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/multiply.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/multiply.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 expressed_type>
+struct Multiply : std::multiplies<void> {};
+
+template <>
+struct Multiply<DataType::kI1> {
+  template <class T>
+  T operator()(const T& lhs, const T& rhs) const {
+    return static_cast<T>(lhs && rhs);
+  }
+};
+
+MultiplyOp Create(MultiplyOp::Attributes) { return {}; }
+
+absl::Status Prepare(MultiplyOp& 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("multiply"), lhs, IsBoolTensor, IsIntTensor,
+                          IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("multiply"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("multiply"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(MultiplyOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  if (IsBoolTensor(lhs)) {
+    detail::EvaluateNoQuantization<DataType::kI1>(Multiply<DataType::kI1>(),
+                                                  lhs, rhs, output);
+    return absl::OkStatus();
+  } else if (IsIntTensor(lhs) || IsFloatTensor(lhs)) {
+    // Note: all the arithmetic types share the same implementation.
+    Multiply<DataType::kF32> multiply;
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       lhs.tensor_element_type(), multiply, lhs, rhs, output);
+  } else if (IsQuantizedPerTensorTensor(lhs)) {
+    Multiply<DataType::kF32> multiply;
+    DISPATCH_QUANTIZED(detail::DequantizeOpQuantizePerTensor,
+                       lhs.quantized_per_tensor_element_type().StorageType(),
+                       lhs.quantized_per_tensor_element_type().ExpressedType(),
+                       multiply, lhs, rhs, output)
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.multiply: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/multiply_test.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 "tensorflow/lite/experimental/shlo/ops/multiply.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/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<MultiplyOp> {
+  static std::string Get() { return "Multiply"; }
+};
+
+template <DataType expressed_type>
+struct Multiply : std::multiplies<void> {};
+
+template <>
+struct Multiply<DataType::kI1> {
+  template <class T>
+  T operator()(const T& lhs, const T& rhs) const {
+    return static_cast<T>(lhs && rhs);
+  }
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Multiply,
+                               BinaryElementwiseOpShapePropagationTest,
+                               MultiplyOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    MultiplyOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes,
+                            BaselineConstraintFloatTypes,
+                            BaselineConstraintQuantizedPerTensorTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Multiply, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<MultiplyOp, PerAxisQuantizedTestTypes>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Multiply, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using ArithmeticTypes = ConcatTypes<BoolTestType, ArithmeticTestTypes>;
+
+template <class T>
+struct MultiplyTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(MultiplyTest, ArithmeticTypes, TestParamNames);
+
+TYPED_TEST(MultiplyTest, ArithmeticTestTypesTensorsWork) {
+  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(),
+                    Multiply<TypeParam::kStorage>());
+
+  auto op = Create(MultiplyOp::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 QuantizedMultiplyTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedMultiplyTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedMultiplyTest, PerTensorWorks) {
+  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 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 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 =
+            Multiply<TypeParam::kExpressed>()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(MultiplyOp::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/negate.cc

@@ -0,0 +1,58 @@
+/* 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/negate.h"
+
+#include <functional>
+
+#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 Negate : std::negate<void> {};
+
+NegateOp Create(NegateOp::Attributes) { return {}; }
+
+absl::Status Prepare(NegateOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(CheckCtx("negate"), input,
+                                               IsSignedIntTensor, IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("negate"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(NegateOp& op, const Tensor& input, Tensor& output) {
+  Negate negate;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), negate,
+        input, output)
+  } else if (IsSignedIntTensor(input) || IsFloatTensor(input)) {
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       input.tensor_element_type(), negate, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.negate: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/negate_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/negate.h"
+
+#include <functional>
+#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;
+using testing::NanSensitiveFloatEq;
+using testing::Pointwise;
+
+namespace shlo_ref {
+
+template <>
+struct ParamName<NegateOp> {
+  static std::string Get() { return "Negate"; }
+};
+
+namespace {
+
+struct Negate : std::negate<void> {
+} negate_ref;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Negate, UnaryElementwiseOpShapePropagationTest,
+                               NegateOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Negate, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<NegateOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<NegateOp, ConcatTypes<BoolTestType, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Negate, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct NegateTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(NegateTest, ArithmeticTestTypes, TestParamNames);
+
+TYPED_TEST(NegateTest, 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(), negate_ref);
+
+  auto op = Create(NegateOp::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 QuantizedNegateTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedNegateTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedNegateTest, 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 = negate_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(NegateOp::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/not.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/not.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 Not {
+  template <class T>
+  T operator()(T v) const {
+    return static_cast<T>(~v);
+  }
+};
+
+template <>
+bool Not::operator()(bool v) const {
+  return !v;
+}
+
+NotOp Create(NotOp::Attributes) { return {}; }
+
+absl::Status Prepare(NotOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("not"), input, IsBoolTensor, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("not"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(NotOp& op, const Tensor& input, Tensor& output) {
+  Not not_func;
+  if (IsIntTensor(input) || IsBoolTensor(input)) {
+    DISPATCH_BOOL_INT(detail::EvaluateNoQuantization,
+                      input.tensor_element_type(), not_func, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.not: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

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

@@ -0,0 +1,103 @@
+/* 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/not.h"
+
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/experimental/shlo/data_type.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<NotOp> {
+  static std::string Get() { return "Not"; }
+};
+
+template <>
+struct SupportedOpDataType<NotOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+struct Not {
+  template <class T>
+  T operator()(T v) const {
+    return ~v;
+  }
+} not_ref;
+
+template <>
+bool Not::operator()(bool v) const {
+  return !v;
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Not, UnaryElementwiseOpShapePropagationTest,
+                               NotOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Not, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    BaselineMismatchSignedIntegerTypes<NotOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<NotOp, ConcatTypes<FloatTestTypes, PerTensorQuantizedTestTypes,
+                                   PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Not, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct BoolAndIntNotTest : ::testing::Test {};
+
+using SupportedTypes = ConcatTypes<BoolTestType, IntTestTypes>;
+
+TYPED_TEST_SUITE(BoolAndIntNotTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(BoolAndIntNotTest, BoolAndIntTensorsWork) {
+  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(), not_ref);
+
+  auto op = Create(NotOp::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/or.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 or
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/ops/or.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 Or : std::bit_or<void> {};
+
+template <>
+struct Or<DataType::kI1> : std::logical_or<void> {};
+
+OrOp Create(OrOp::Attributes) { return {}; }
+
+absl::Status Prepare(OrOp& 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("or"), lhs, IsBoolTensor, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("or"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("or"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(OrOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  if (IsIntTensor(lhs)) {
+    // Note: all the integer types share the same implementation.
+    Or<DataType::kSI32> or_func;
+    DISPATCH_INT(detail::EvaluateNoQuantization, lhs.tensor_element_type(),
+                 or_func, lhs, rhs, output);
+  } else if (IsBoolTensor(lhs)) {
+    Or<DataType::kI1> or_func;
+    detail::EvaluateNoQuantization<DataType::kI1>(or_func, lhs, rhs, output);
+    return absl::OkStatus();
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.or: Unsupported tensor type in Evaluate.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/or.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 or
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_OR_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_OR_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct OrOp {
+  struct Attributes {};
+};
+
+OrOp Create(OrOp::Attributes);
+absl::Status Prepare(OrOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output);
+absl::Status Evaluate(OrOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_OR_H_

tensorflow/lite/experimental/shlo/ops/or_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 or
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/ops/or.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<OrOp> {
+  static std::string Get() { return "Or"; }
+};
+
+template <DataType>
+struct Or : std::bit_or<void> {};
+
+template <>
+struct Or<DataType::kI1> : std::logical_or<void> {};
+
+template <>
+struct SupportedOpDataType<OrOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Or, BinaryElementwiseOpShapePropagationTest,
+                               OrOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    OrOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Or, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<OrOp, ConcatTypes<FloatTestTypes, PerTensorQuantizedTestTypes,
+                                  PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Or, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, IntTestTypes>;
+
+template <class T>
+struct OrTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(OrTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(OrTest, 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(),
+                    Or<TypeParam::kStorage>());
+
+  auto op = Create(OrOp::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/popcnt.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/popcnt.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 Popcnt {
+  template <class T>
+  T operator()(T v) const {
+    if constexpr (std::is_same_v<I4, T>) {
+      return I4(absl::popcount(static_cast<uint8_t>(v & 0xf)));
+    } else {
+      return absl::popcount(static_cast<std::make_unsigned_t<T>>(v));
+    }
+  }
+};
+
+PopcntOp Create(PopcntOp::Attributes) { return {}; }
+
+absl::Status Prepare(PopcntOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSupportedTypes(CheckCtx("popcnt"), input, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("popcnt"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(PopcntOp& op, const Tensor& input, Tensor& output) {
+  Popcnt popcnt;
+  if (IsIntTensor(input)) {
+    DISPATCH_INT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                 popcnt, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.popcnt: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

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

@@ -0,0 +1,127 @@
+/* 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/popcnt.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<PopcntOp> {
+  static std::string Get() { return "Popcnt"; }
+};
+
+template <>
+struct SupportedOpDataType<PopcntOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+struct Popcnt {
+  template <class T>
+  T operator()(T v) const {
+    if constexpr (std::is_same_v<I4, T>) {
+      return I4(absl::popcount(static_cast<uint8_t>(v & 0xf)));
+    } else {
+      return absl::popcount(static_cast<std::make_unsigned_t<T>>(v));
+    }
+  }
+} popcnt_ref;
+
+using PopcntTypes = ::testing::Types<int32_t, int16_t, int8_t, I4>;
+
+template <class T>
+struct PopcntFunctorTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(PopcntFunctorTest, PopcntTypes);
+
+TYPED_TEST(PopcntFunctorTest, GivesCorrectResults) {
+  int64_t bit_count = 8 * sizeof(TypeParam);
+  if constexpr (std::is_same_v<I4, TypeParam>) {
+    bit_count = 4;
+  }
+  EXPECT_EQ(popcnt_ref(std::numeric_limits<TypeParam>::lowest()), 1);
+  EXPECT_EQ(popcnt_ref(static_cast<TypeParam>(-1)), bit_count);
+  EXPECT_EQ(popcnt_ref(static_cast<TypeParam>(0)), 0);
+  EXPECT_EQ(popcnt_ref(static_cast<TypeParam>(1)), 1);
+  EXPECT_EQ(popcnt_ref(static_cast<TypeParam>(2)), 1);
+  EXPECT_EQ(popcnt_ref(static_cast<TypeParam>(3)), 2);
+  EXPECT_EQ(popcnt_ref(std::numeric_limits<TypeParam>::max()), bit_count - 1);
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Popcnt, UnaryElementwiseOpShapePropagationTest,
+                               PopcntOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Popcnt, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    BaselineMismatchSignedIntegerTypes<PopcntOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<PopcntOp, ConcatTypes<BoolTestType, FloatTestTypes,
+                                      PerTensorQuantizedTestTypes,
+                                      PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Popcnt, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct PopcntTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(PopcntTest, IntTestTypes, TestParamNames);
+
+TYPED_TEST(PopcntTest, 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(), popcnt_ref);
+
+  auto op = Create(PopcntOp::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/sign.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/sign.h"
+
+#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 Sign {
+  template <class T>
+  T operator()(T v) const {
+    constexpr T one = static_cast<T>(1);
+    constexpr T minus_one = static_cast<T>(-1);
+    constexpr T zero = static_cast<T>(0);
+    return v < zero ? minus_one : (v > zero ? one : v);
+  }
+};
+
+template <>
+F16 Sign::operator()(F16 v) const {
+  return static_cast<F16>(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 Sign::operator()(BF16 v) const {
+  return static_cast<BF16>(operator()(static_cast<float>(v)));
+}
+
+SignOp Create(SignOp::Attributes) { return {}; }
+
+absl::Status Prepare(SignOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(CheckCtx("sign"), input,
+                                               IsSignedIntTensor, IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("sign"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(SignOp& op, const Tensor& input, Tensor& output) {
+  Sign sign;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), sign, input,
+        output)
+  } else if (IsSignedIntTensor(input) || IsFloatTensor(input)) {
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       input.tensor_element_type(), sign, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.sign: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/sign_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/sign.h"
+
+#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<SignOp> {
+  static std::string Get() { return "Sign"; }
+};
+
+namespace {
+
+struct Sign {
+  template <class T>
+  T operator()(T v) const {
+    constexpr T one = static_cast<T>(1);
+    constexpr T minus_one = static_cast<T>(-1);
+    constexpr T zero = static_cast<T>(0);
+    return v < zero ? minus_one : (v > zero ? one : v);
+  }
+} sign_ref;
+
+template <>
+F16 Sign::operator()(F16 v) const {
+  return static_cast<F16>(operator()(static_cast<float>(v)));
+}
+
+template <>
+BF16 Sign::operator()(BF16 v) const {
+  return static_cast<BF16>(operator()(static_cast<float>(v)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sign, UnaryElementwiseOpShapePropagationTest,
+                               SignOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Sign, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<SignOp>, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<SignOp, ConcatTypes<BoolTestType, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sign, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct SignTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(SignTest, ArithmeticTestTypes, TestParamNames);
+
+TYPED_TEST(SignTest, 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(), sign_ref);
+
+  auto op = Create(SignOp::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 QuantizedSignTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedSignTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedSignTest, 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 = sign_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(SignOp::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/sine.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/sine.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/data_type.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 Sine {
+  template <class T>
+  T operator()(T v) const {
+    return std::sin(v);
+  }
+};
+
+template <>
+F16 Sine::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Sine::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+SineOp Create(SineOp::Attributes) { return {}; }
+
+absl::Status Prepare(SineOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("sine"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("sine"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(SineOp& op, const Tensor& input, Tensor& output) {
+  Sine sine;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), sine, input,
+        output)
+  } else if (!input.IsQuantized() && IsFloat(input.StorageType())) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   sine, input, output);
+  }
+  return absl::FailedPreconditionError("Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/sine_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/sine.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<SineOp> {
+  static std::string Get() { return "Sine"; }
+};
+
+namespace {
+
+struct Sine {
+  template <class T>
+  T operator()(T v) const {
+    return std::sin(v);
+  }
+} sine_ref;
+
+template <>
+F16 Sine::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Sine::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sine, UnaryElementwiseOpShapePropagationTest,
+                               SineOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Sine, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<SineOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    SineOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sine, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct FloatSineTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(FloatSineTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(FloatSineTest, 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());
+  const TensorType tensor_type =
+      TensorType{.shape = shape, .element_type = TypeParam::kStorage};
+  Tensor input_tensor{.type = tensor_type, .data = input_data.data()};
+  Tensor output_tensor{.type = tensor_type, .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(input_data, expected_data.begin(), sine_ref);
+
+  auto op = Create(SineOp::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 QuantizedSineTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedSineTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedSineTest, 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);
+  Vector<StorageT> output_data(shape.NumElements());
+  const QuantizedPerTensorTensorType tensor_type = {
+      .shape = shape,
+      .element_type = QuantizedElementTypePerTensor(
+          TypeParam::kStorage, zero_point, TypeParam::kExpressed, scale)};
+  Tensor input_tensor{.type = tensor_type, .data = input_data.data()};
+  Tensor output_tensor{.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 = sine_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(SineOp::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/sqrt.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/sqrt.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 Sqrt {
+  template <class T>
+  T operator()(T v) const {
+    return std::sqrt(v);
+  }
+};
+
+template <>
+F16 Sqrt::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Sqrt::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+SqrtOp Create(SqrtOp::Attributes) { return {}; }
+
+absl::Status Prepare(SqrtOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("sqrt"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("sqrt"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(SqrtOp& op, const Tensor& input, Tensor& output) {
+  Sqrt sqrt;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), sqrt, input,
+        output)
+  } else if (IsFloatTensor(input)) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   sqrt, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.sqrt: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/sqrt_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/sqrt.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<SqrtOp> {
+  static std::string Get() { return "Sqrt"; }
+};
+
+namespace {
+
+struct Sqrt {
+  template <class T>
+  T operator()(T v) const {
+    return std::sqrt(v);
+  }
+} sqrt_ref;
+
+template <>
+F16 Sqrt::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Sqrt::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sqrt, UnaryElementwiseOpShapePropagationTest,
+                               SqrtOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Sqrt, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<SqrtOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    SqrtOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Sqrt, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct FloatSqrtTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(FloatSqrtTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(FloatSqrtTest, 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));
+  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(), sqrt_ref);
+
+  auto op = Create(SqrtOp::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 QuantizedSqrtTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedSqrtTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedSqrtTest, 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>(4);
+  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 = sqrt_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(SqrtOp::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/subtract.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/subtract.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 Subtract : std::minus<void> {};
+
+SubtractOp Create(SubtractOp::Attributes) { return {}; }
+
+absl::Status Prepare(SubtractOp& 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("subtract"), lhs,
+                                               IsIntTensor, IsFloatTensor,
+                                               IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("subtract"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("subtract"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(SubtractOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  Subtract subtract;
+  if (IsIntTensor(lhs) || IsFloatTensor(lhs)) {
+    // Note: all the arithmetic types share the same implementation.
+    DISPATCH_INT_FLOAT(detail::EvaluateNoQuantization,
+                       lhs.tensor_element_type(), subtract, 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(),
+                       subtract, lhs, rhs, output)
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.subtract: Unsupported tensor type.");
+}
+
+}  // namespace shlo_ref

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

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

@@ -0,0 +1,151 @@
+/* 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/subtract.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<SubtractOp> {
+  static std::string Get() { return "Subtract"; }
+};
+
+struct Subtract : std::minus<void> {};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Subtract,
+                               BinaryElementwiseOpShapePropagationTest,
+                               SubtractOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    SubtractOp,
+    ConcatTypes<BaselineConstraintIntTypes, BaselineConstraintFloatTypes,
+                BaselineConstraintQuantizedPerTensorTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Subtract, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<SubtractOp,
+                ConcatTypes<BoolTestType, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Subtract, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using ArithmeticTypes = ConcatTypes<ArithmeticTestTypes>;
+
+template <class T>
+struct SubtractTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(SubtractTest, ArithmeticTypes, TestParamNames);
+
+TYPED_TEST(SubtractTest, 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=*/-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(), Subtract());
+
+  auto op = Create(SubtractOp::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 QuantizedSubtractTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedSubtractTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedSubtractTest, 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=*/-5, /*max=*/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 =
+            Subtract()(dequantized_lhs, dequantized_rhs);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(SubtractOp::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/tanh.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/tanh.h"
+
+#include <cmath>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/bf16.h"
+#include "tensorflow/lite/experimental/shlo/data_type.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 Tanh {
+  template <class T>
+  T operator()(T v) const {
+    return std::tanh(v);
+  }
+};
+
+template <>
+F16 Tanh::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Tanh::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+TanhOp Create(TanhOp::Attributes) { return {}; }
+
+absl::Status Prepare(TanhOp& op, const Tensor& input, Tensor& output) {
+  SHLO_REF_RETURN_ON_ERROR(Propagate(input.shape(), output.shape()));
+  SHLO_REF_RETURN_ON_ERROR(CheckSupportedTypes(
+      CheckCtx("tanh"), input, IsFloatTensor, IsQuantizedPerTensorTensor));
+  SHLO_REF_RETURN_ON_ERROR(
+      CheckSameBaselineType(CheckCtx("tanh"), input, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(TanhOp& op, const Tensor& input, Tensor& output) {
+  Tanh tanh;
+  if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), tanh, input,
+        output)
+  } else if (!input.IsQuantized() && IsFloat(input.StorageType())) {
+    DISPATCH_FLOAT(detail::EvaluateNoQuantization, input.tensor_element_type(),
+                   tanh, input, output);
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.tanh: Unsupported tensor type.");
+}
+
+};  // namespace shlo_ref

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

tensorflow/lite/experimental/shlo/ops/tanh_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/tanh.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<TanhOp> {
+  static std::string Get() { return "Tanh"; }
+};
+
+namespace {
+
+struct Tanh {
+  template <class T>
+  T operator()(T v) const {
+    return std::tanh(v);
+  }
+} tanh_ref;
+
+template <>
+F16 Tanh::operator()<F16>(F16 val) const {
+  return F16(operator()(static_cast<float>(val)));
+}
+
+template <>
+BF16 Tanh::operator()<BF16>(BF16 val) const {
+  return BF16(operator()(static_cast<float>(val)));
+}
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Tanh, UnaryElementwiseOpShapePropagationTest,
+                               TanhOp, TestParamNames);
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Tanh, UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    UnaryElementwiseConstraint1Types<TanhOp>, TestParamNames);
+
+using UnsupportedTypes = WithOpTypes<
+    TanhOp, ConcatTypes<BoolTestType, IntTestTypes, PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Tanh, UnaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+template <class T>
+struct FloatTanhTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(FloatTanhTest, FloatTestTypes, TestParamNames);
+
+TYPED_TEST(FloatTanhTest, 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(), tanh_ref);
+
+  auto op = Create(TanhOp::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 QuantizedTanhTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedTanhTest, QuantizedTestTypes, TestParamNames);
+
+TYPED_TEST(QuantizedTanhTest, 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 = tanh_ref(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(TanhOp::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/test_util.h

@@ -0,0 +1,442 @@
+/* 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_TEST_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_TEST_UTIL_H_
+
+#include <cstdint>
+#include <random>
+#include <string>
+#include <tuple>
+#include <type_traits>
+
+#include <gtest/gtest.h>
+#include "absl/algorithm/container.h"
+#include "absl/container/inlined_vector.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/i4.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 {
+
+// We use a vector class that is different from std::vector to have a consistent
+// API when dealing with bool tensors.
+template <class T>
+using Vector = absl::InlinedVector<T, 1>;
+
+// Helper for UniformDistribution.
+template <DataType storage_type, typename = void>
+struct UniformDistributionImpl;
+
+template <>
+struct UniformDistributionImpl<DataType::kI1, void>
+    : std::uniform_int_distribution<int32_t> {
+  using std::uniform_int_distribution<int32_t>::uniform_int_distribution;
+};
+
+template <>
+struct UniformDistributionImpl<DataType::kSI4, void>
+    : std::uniform_int_distribution<int8_t> {
+  using std::uniform_int_distribution<int8_t>::uniform_int_distribution;
+};
+
+template <DataType storage_type>
+struct UniformDistributionImpl<storage_type,
+                               std::enable_if_t<IsInteger(storage_type)>>
+    : std::uniform_int_distribution<typename Storage<storage_type>::Type> {
+  using std::uniform_int_distribution<
+      typename Storage<storage_type>::Type>::uniform_int_distribution;
+};
+
+template <DataType storage_type>
+struct UniformDistributionImpl<storage_type,
+                               std::enable_if_t<IsFloat(storage_type)>>
+    : std::uniform_real_distribution<float> {
+  using std::uniform_real_distribution<float>::uniform_real_distribution;
+};
+
+// Helps creating a uniform distribution for the given data type.
+template <DataType storage_type>
+using UniformDistribution = UniformDistributionImpl<storage_type>;
+
+// Returns a vector filled with random data according to the set distribution.
+template <DataType storage_type,
+          template <DataType> class Distribution = UniformDistribution,
+          class MinT = StorageType<storage_type>,
+          class MaxT = StorageType<storage_type>,
+          class Config = Storage<storage_type>>
+Vector<typename Config::Type> RandomBuffer(const Shape& shape,
+                                           const MinT min = Config::kMinValue,
+                                           const MaxT max = Config::kMaxValue) {
+  using StorageT = StorageType<storage_type>;
+  const StorageT min_val =
+      min > Config::kMinValue ? static_cast<StorageT>(min) : Config::kMinValue;
+  const StorageT max_val =
+      max < Config::kMaxValue ? static_cast<StorageT>(max) : Config::kMaxValue;
+  Vector<typename Config::Type> vec(shape.NumElements());
+  std::random_device rd;
+  if constexpr (std::is_same_v<I4, StorageT>) {
+    Distribution<DataType::kSI8> dist(min_val, max_val);
+    absl::c_generate(vec, [&] { return static_cast<StorageT>(dist(rd)); });
+  } else {
+    Distribution<storage_type> dist(min_val, max_val);
+    absl::c_generate(vec, [&] { return dist(rd); });
+  }
+  return vec;
+}
+
+// Returns a vector filled with incremental value. The values wrap around
+// according to the storage type range.
+template <DataType storage_type, class StartT = StorageType<storage_type>,
+          class MinT = StorageType<storage_type>,
+          class MaxT = StorageType<storage_type>,
+          class Config = Storage<storage_type>>
+Vector<typename Config::Type> IotaBuffer(const Shape& shape,
+                                         const StartT start = Config::kMinValue,
+                                         const MinT min = Config::kMinValue,
+                                         const MaxT max = Config::kMaxValue) {
+  using StorageT = StorageType<storage_type>;
+  const StorageT min_val =
+      min > Config::kMinValue ? static_cast<StorageT>(min) : Config::kMinValue;
+  const StorageT max_val =
+      max < Config::kMaxValue ? static_cast<StorageT>(max) : Config::kMaxValue;
+  Vector<typename Config::Type> vec(shape.NumElements());
+  StorageT v = start >= min_val ? static_cast<StorageT>(start) : min_val;
+  v = v <= max_val ? v : min_val;
+  for (auto& e : vec) {
+    e = v;
+    if (v >= max_val) {
+      v = min_val;
+    } else {
+      ++v;
+    }
+  }
+  return vec;
+}
+
+// Typed test parameter type.
+template <DataType... Types>
+struct TestParam;
+
+// Typed test parameter specialization for non quantized tensors.
+template <DataType storage_type>
+struct TestParam<storage_type> {
+  static constexpr DataType kStorage = storage_type;
+  using StorageT = StorageType<storage_type>;
+};
+
+// Typed test parameter specialization for quantized tensors.
+template <DataType storage_type, DataType expressed_type>
+struct TestParam<storage_type, expressed_type> {
+  static constexpr DataType kStorage = storage_type;
+  static constexpr DataType kExpressed = expressed_type;
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+};
+
+// Typed test parameter tag to ask for a per-tensor quantized tensor.
+//
+// TestParamT should be a `TestParam<storage_type, expressed_type>`.
+template <class TestParamT>
+struct PerTensor {
+  using Param = TestParamT;
+};
+
+// Typed test parameter tag to ask for a per-channel quantized tensor.
+//
+// TestParamT should be a `TestParam<storage_type, expressed_type>`.
+template <class TestParamT, Axis kAxis = 0>
+struct PerAxis {
+  using Param = TestParamT;
+  static constexpr Axis axis = kAxis;
+};
+
+// Helps getting a human readable typed test parameter name.
+template <class T>
+struct ParamName;
+
+template <DataType T, DataType... Ts>
+struct ParamName<TestParam<T, Ts...>> {
+  static std::string Get() {
+    std::string name = std::string("") + ToString(T);
+    ((name += std::string("_") + ToString(Ts)), ...);
+    return name;
+  }
+};
+
+template <DataType T, DataType... Ts>
+struct ParamName<PerTensor<TestParam<T, Ts...>>> {
+  static std::string Get() {
+    std::string name = std::string("PerTensor[") + ToString(T);
+    ((name += std::string("_") + ToString(Ts)), ...);
+    return name + "]";
+  }
+};
+
+template <DataType T, DataType... Ts, Axis axis>
+struct ParamName<PerAxis<TestParam<T, Ts...>, axis>> {
+  static std::string Get() {
+    std::string name = std::string("PerAxis[") + ToString(T);
+    ((name += std::string("_") + ToString(Ts)), ...);
+    return name + ":" + std::to_string(axis) + "]";
+  }
+};
+
+template <class TestParamT, class... TestParamTs>
+struct ParamName<std::tuple<TestParamT, TestParamTs...>> {
+  static std::string Get() {
+    std::string name = ParamName<TestParamT>::Get();
+    ((name += std::string(":") + ParamName<TestParamTs>::Get()), ...);
+    return name;
+  }
+};
+
+// Allows GTest to print a human readable version of the typed test parameters.
+class TestParamNames {
+ public:
+  template <class T>
+  static std::string GetName(int) {
+    return ParamName<T>::Get();
+  }
+};
+
+// Applies the F template to the given testing::Types list.
+template <template <class> class F, class T>
+struct Map;
+
+template <template <class> class F, class... Ts>
+struct Map<F, ::testing::Types<Ts...>> {
+  using Types = ::testing::Types<F<Ts>...>;
+};
+
+template <template <class> class F, class T>
+using MapTypes = typename Map<F, T>::Types;
+
+// Concatenates testing::Types lists.
+template <class... Ts>
+struct Concat;
+
+template <class... Ts>
+struct Concat<::testing::Types<Ts...>> {
+  using Types = ::testing::Types<Ts...>;
+};
+
+template <class... Ts, class... Us, class... ExtraTypes>
+struct Concat<::testing::Types<Ts...>, ::testing::Types<Us...>, ExtraTypes...> {
+  using Types =
+      typename Concat<::testing::Types<Ts..., Us...>, ExtraTypes...>::Types;
+};
+
+template <class... Ts>
+using ConcatTypes = typename Concat<Ts...>::Types;
+
+// Transforms a list of types into a list of tuple<Op, type>.
+template <class Op, class T>
+struct WithOp;
+
+template <class Op, class... Ts>
+struct WithOp<Op, ::testing::Types<Ts...>> {
+  using Types = ::testing::Types<std::tuple<Op, Ts>...>;
+};
+
+template <class Op, class T>
+using WithOpTypes = typename WithOp<Op, T>::Types;
+
+// Helps generating a cross-product of lists.
+template <class Accu, class... Lists>
+struct CrossProductImpl;
+
+template <class... AccuTs, class... Ts, class... Lists>
+struct CrossProductImpl<::testing::Types<AccuTs...>, ::testing::Types<Ts...>,
+                        Lists...> {
+  using Types =
+      ConcatTypes<typename CrossProductImpl<::testing::Types<AccuTs..., Ts>,
+                                            Lists...>::Types...>;
+};
+
+template <class... AccuTs>
+struct CrossProductImpl<::testing::Types<AccuTs...>> {
+  using Types = ::testing::Types<::testing::Types<AccuTs...>>;
+};
+
+// Generates a cross-product of lists.
+template <class... Lists>
+struct CrossProduct {
+  using Types = typename CrossProductImpl<::testing::Types<>, Lists...>::Types;
+};
+
+template <class... Lists>
+using CrossProductTypes = typename CrossProduct<Lists...>::Types;
+
+static_assert(
+    std::is_same_v<
+        CrossProductTypes<::testing::Types<int, float>,
+                          ::testing::Types<char, double>>,
+        ::testing::Types<
+            ::testing::Types<int, char>, ::testing::Types<int, double>,
+            ::testing::Types<float, char>, ::testing::Types<float, double>>>);
+
+static_assert(
+    std::is_same_v<
+        CrossProductTypes<::testing::Types<int>, ::testing::Types<char, double>,
+                          ::testing::Types<float>>,
+        ::testing::Types<::testing::Types<int, char, float>,
+                         ::testing::Types<int, double, float>>>);
+
+// Filters out the types that don't satisfy the predicate.
+template <template <class...> class Predicate, class List>
+struct Filter;
+
+template <template <class...> class Predicate, class... Ts>
+struct Filter<Predicate, ::testing::Types<Ts...>> {
+  using Type =
+      ConcatTypes<std::conditional_t<Predicate<Ts>::value, ::testing::Types<Ts>,
+                                     ::testing::Types<>>...>;
+};
+
+template <template <class...> class Predicate, class List>
+using FilterTypes = typename Filter<Predicate, List>::Type;
+
+static_assert(std::is_same_v<
+              FilterTypes<std::is_integral, ::testing::Types<int, char, float>>,
+              ::testing::Types<int, char>>);
+
+// Checks if all given types are the same.
+template <class T, class... Ts>
+struct SameTypes : std::bool_constant<(std::is_same_v<T, Ts> && ...)> {};
+
+// Checks if all types in the testing::Types list are the same.
+template <class T, class... Ts>
+struct SameTypes<::testing::Types<T, Ts...>> : SameTypes<T, Ts...> {};
+
+// Provides a new predicate that negates the given one.
+template <template <class...> class Pred>
+struct NegatePred {
+  template <class... Ts>
+  using Predicate = std::negation<Pred<Ts...>>;
+};
+
+// Use this with TYPED_TEST_SUITE for boolean testing.
+using BoolTestType = ::testing::Types<TestParam<DataType::kI1>>;
+
+// Use this with TYPED_TEST_SUITE for non quantized integer testing.
+using IntTestTypes =
+    ::testing::Types<TestParam<DataType::kSI4>, TestParam<DataType::kSI8>,
+                     TestParam<DataType::kSI16>, TestParam<DataType::kSI32>>;
+
+// Use this with TYPED_TEST_SUITE for non quantized floating point testing.
+using FloatTestTypes =
+    ::testing::Types<TestParam<DataType::kBF16>, TestParam<DataType::kF16>,
+                     TestParam<DataType::kF32>>;
+
+// Use this with TYPED_TEST_SUITE for non quantized testing.
+using ArithmeticTestTypes = ConcatTypes<IntTestTypes, FloatTestTypes>;
+
+// Use this with TYPED_TEST_SUITE for unspecified quantized testing.
+using QuantizedTestTypes =
+    ::testing::Types<TestParam<DataType::kSI4, DataType::kF32>,
+                     TestParam<DataType::kSI8, DataType::kF32>,
+                     TestParam<DataType::kSI16, DataType::kF32>,
+                     TestParam<DataType::kSI4, DataType::kBF16>,
+                     TestParam<DataType::kSI8, DataType::kBF16>,
+                     TestParam<DataType::kSI4, DataType::kF16>,
+                     TestParam<DataType::kSI8, DataType::kF16>>;
+
+// Use this with TYPED_TEST_SUITE for quantized per tensor testing.
+using PerTensorQuantizedTestTypes = MapTypes<PerTensor, QuantizedTestTypes>;
+
+template <class T>
+using PerAxis0 = PerAxis<T, 0>;
+
+// Use this with TYPED_TEST_SUITE for quantized per axis testing.
+using PerAxisQuantizedTestTypes = MapTypes<PerAxis0, QuantizedTestTypes>;
+
+// Customization point for generic tests that need to create a supported tensor
+// for an op but that don't care what that type is.
+//
+// Specialize this in the test file if F32 isn't supported by the op under test.
+template <class Op>
+struct SupportedOpDataType {
+  static constexpr DataType kStorageType = DataType::kF32;
+};
+
+// Customization point for generic tests that need to create a supported output
+// tensor for an op but that don't care what that type is.
+//
+// Specialize this in the test file if `SupportedOpDataType<Op>::kStorageType`
+// isn't supported by the op under test.
+template <class Op>
+struct SupportedOpOutputDataType {
+  static constexpr DataType kStorageType =
+      SupportedOpDataType<Op>::kStorageType;
+};
+
+// Customization point for generic tests that need a valid attribute
+// configuration to create an op but that don't care what that configuration is.
+//
+// Specialize this in the test file if F32 isn't supported by the op under test.
+template <class Op>
+struct SupportedOpAttributes {
+  static typename Op::Attributes Get() { return {}; };
+};
+
+// Builds a TensorType object and returns it in a variant that can be passed to
+// a tensor.
+template <DataType storage_type>
+TensorTypeVariant TensorTypeFor(TestParam<storage_type>, const Shape& shape) {
+  return TensorType{.shape = shape, .element_type = storage_type};
+}
+
+// Builds a per tensor QuantizedTensorType object and returns it in a variant
+// that can be passed to a tensor.
+//
+// WARNING: the scale and zero point are randomly generated:
+//   - scale is in [0.5, 1.5]
+//   - zero_point is in [-5, 5]
+template <DataType storage_type, DataType expressed_type>
+TensorTypeVariant TensorTypeFor(
+    PerTensor<TestParam<storage_type, expressed_type>>, const Shape& shape) {
+  std::random_device rd;
+  UniformDistribution<expressed_type> expressed_dist(0.5, 1.5);
+  UniformDistribution<storage_type> storage_dist(-5, 5);
+  StorageType<expressed_type> scale =
+      static_cast<StorageType<expressed_type>>(expressed_dist(rd));
+  StorageType<storage_type> zero_point =
+      StorageType<storage_type>(storage_dist(rd));
+  return QuantizedPerTensorTensorType{
+      .shape = shape,
+      .element_type = QuantizedElementTypePerTensor(storage_type, zero_point,
+                                                    expressed_type, scale)};
+}
+
+// Builds a per axis QuantizedTensorType object and returns it in a variant
+// that can be passed to a tensor.
+//
+// WARNING: scales and zero points are unspecified and may be empty.
+template <DataType storage_type, DataType expressed_type, Axis axis>
+TensorTypeVariant TensorTypeFor(
+    PerAxis<TestParam<storage_type, expressed_type>, axis>,
+    const Shape& shape) {
+  return QuantizedPerAxisTensorType{
+      .shape = shape,
+      .element_type = QuantizedElementTypePerAxis(storage_type, {},
+                                                  expressed_type, {}, axis)};
+}
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_TEST_UTIL_H_

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

@@ -0,0 +1,192 @@
+/* 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_UNARY_ELEMENTWISE_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UNARY_ELEMENTWISE_H_
+
+#include <cstddef>
+
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/dispatch.h"
+#include "tensorflow/lite/experimental/shlo/ops/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/tensor.h"
+
+namespace shlo_ref {
+
+namespace detail {
+
+template <typename StorageT, typename ExpressedT, typename F>
+void DequantizeOpQuantizePerAxisImpl(
+    F& op, const Shape& shape, const Axis quantization_dimension,
+    const StorageT quantization_min, const StorageT quantization_max,
+    const absl::Span<const StorageT> input_zero_points,
+    const absl::Span<const ExpressedT> input_scales,
+    const absl::Span<const StorageT> output_zero_points,
+    const absl::Span<const ExpressedT> output_scales, const Strides& strides,
+    const StorageT* input_data, StorageT* output_data, const size_t depth,
+    size_t quantization_index) {
+  const DimensionSize dim = shape.Dim(depth);
+  if (depth + 1 >= shape.Rank()) {
+    for (DimensionSize i = 0; i < dim; ++i) {
+      if (depth == quantization_dimension) {
+        quantization_index = i;
+      }
+      const ExpressedT dequantized_input =
+          Dequantize(*input_data, input_zero_points[quantization_index],
+                     input_scales[quantization_index]);
+      const ExpressedT dequantized_res = op(dequantized_input);
+      *output_data = Quantize<StorageT, ExpressedT>(
+          dequantized_res, output_zero_points[quantization_index],
+          static_cast<ExpressedT>(1) / output_scales[quantization_index],
+          quantization_min, quantization_max);
+      output_data += strides[depth];
+      input_data += strides[depth];
+    }
+  } else {
+    for (DimensionSize i = 0; i < dim; ++i) {
+      if (depth == quantization_dimension) {
+        quantization_index = i;
+      }
+      DequantizeOpQuantizePerAxisImpl(
+          op, shape, quantization_dimension, quantization_min, quantization_max,
+          input_zero_points, input_scales, output_zero_points, output_scales,
+          strides, input_data, output_data, depth + 1, quantization_index);
+      output_data += strides[depth];
+      input_data += strides[depth];
+    }
+  }
+}
+
+template <DataType storage_type, DataType expressed_type, typename F>
+void DequantizeOpQuantizePerAxis(F&& func, const Tensor& input,
+                                 Tensor& output) {
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+  const Shape& shape = input.shape();
+  const Axis quantization_dimension =
+      input.quantized_per_axis_element_type().QuantizedDimension();
+  const absl::Span<const StorageT> input_zero_points =
+      input.quantized_per_axis_element_type().ZeroPointsAs<storage_type>();
+  const absl::Span<const ExpressedT> input_scales =
+      input.quantized_per_axis_element_type().ScalesAs<expressed_type>();
+  const absl::Span<const StorageT> output_zero_points =
+      output.quantized_per_axis_element_type().ZeroPointsAs<storage_type>();
+  const absl::Span<const ExpressedT> output_scales =
+      output.quantized_per_axis_element_type().ScalesAs<expressed_type>();
+  const Strides& strides = ComputeStrides(shape);
+  const StorageT* input_data = input.GetDataAs<storage_type>();
+  StorageT* output_data = output.GetDataAs<storage_type>();
+  DequantizeOpQuantizePerAxisImpl(
+      func, shape, quantization_dimension, Storage<storage_type>::kMinValue,
+      Storage<storage_type>::kMaxValue, input_zero_points, input_scales,
+      output_zero_points, output_scales, strides, input_data, output_data,
+      /*depth=*/0, /*quantization_index=*/0);
+}
+
+template <DataType storage_type, DataType expressed_type, typename F>
+void DequantizeOpQuantizePerTensor(F& func, const Tensor& input,
+                                   Tensor& output) {
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+  const DimensionSize num_elements = input.NumElements();
+  const StorageT input_zero_point =
+      input.quantized_per_tensor_element_type().ZeroPointAs<storage_type>();
+  const ExpressedT input_scale =
+      input.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* input_data = input.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, ++input_data, ++output_data) {
+    const ExpressedT dequantized_input =
+        Dequantize(*input_data, input_zero_point, input_scale);
+    const ExpressedT dequantized_res = func(dequantized_input);
+    *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& input, Tensor& output) {
+  absl::c_transform(input.Flat<data_type>(), output.GetDataAs<data_type>(),
+                    static_cast<F&&>(func));
+}
+
+}  // namespace detail
+
+// The following structures and functions are examples to implement unary ops.
+
+template <class F>
+struct UnaryElementwiseOp {
+  struct Attributes {};
+  F func;
+};
+
+// Creates the op structure and initializes the functor if it has a state.
+template <class F>
+UnaryElementwiseOp<F> Create(typename UnaryElementwiseOp<F>::Attributes,
+                             const F& func) {
+  return UnaryElementwiseOp<F>{func};
+}
+
+template <class F>
+UnaryElementwiseOp<F> Create(typename UnaryElementwiseOp<F>::Attributes,
+                             F&& func) {
+  return UnaryElementwiseOp<F>{static_cast<F&&>(func)};
+}
+
+// Checks the op constraints and propagates the output shape if needed.
+template <class F>
+absl::Status Prepare(UnaryElementwiseOp<F>& op, const Tensor& input,
+                     Tensor& output) {
+  return Propagate(input.shape(), output.shape());
+}
+
+// Runs the op over the input tensor.
+template <class F>
+absl::Status Evaluate(UnaryElementwiseOp<F>& op, const Tensor& input,
+                      Tensor& output) {
+  if (input.IsPerAxisQuantized()) {
+    DISPATCH_QUANTIZED(detail::DequantizeOpQuantizePerAxis,
+                       input.quantized_per_axis_element_type().StorageType(),
+                       input.quantized_per_axis_element_type().ExpressedType(),
+                       op.func, input, output);
+  } else if (input.IsPerTensorQuantized()) {
+    DISPATCH_QUANTIZED(
+        detail::DequantizeOpQuantizePerTensor,
+        input.quantized_per_tensor_element_type().StorageType(),
+        input.quantized_per_tensor_element_type().ExpressedType(), op.func,
+        input, output)
+  } else {
+    DISPATCH_BOOL_INT_FLOAT(detail::EvaluateNoQuantization,
+                            input.tensor_element_type(), op.func, input,
+                            output);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UNARY_ELEMENTWISE_H_

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

@@ -0,0 +1,181 @@
+/* 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/unary_elementwise.h"
+
+#include <cstddef>
+#include <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/algorithm/container.h"
+#include "tensorflow/lite/experimental/shlo/data_type.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/status_matcher.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+using testing::ElementsAreArray;
+
+namespace shlo_ref {
+namespace {
+
+struct Abs {
+  template <class T>
+  T operator()(const T& val) {
+    return val < static_cast<T>(0) ? static_cast<T>(-val) : val;
+  }
+};
+
+template <DataType storage_type, DataType expressed_type = DataType::kF32>
+struct TestParam {
+  static constexpr DataType kStorage = storage_type;
+  static constexpr DataType kExpressed = expressed_type;
+  using StorageT = StorageType<storage_type>;
+  using ExpressedT = StorageType<expressed_type>;
+};
+
+template <class T>
+struct UnaryElementWiseTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(UnaryElementWiseTest, ArithmeticTestTypes);
+
+TYPED_TEST(UnaryElementWiseTest, NonQuantizedWithAbs) {
+  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());
+
+  auto op = Create(UnaryElementwiseOp<Abs>::Attributes{}, Abs());
+  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 QuantizedUnaryElementWiseTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(QuantizedUnaryElementWiseTest, QuantizedTestTypes);
+
+TYPED_TEST(QuantizedUnaryElementWiseTest, QuantizedPerTensorWithAbs) {
+  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()(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, static_cast<ExpressedT>(1.) / scale);
+      });
+
+  auto op = Create(UnaryElementwiseOp<Abs>::Attributes{}, Abs());
+  ASSERT_OK(Prepare(op, input_tensor, output_tensor));
+  ASSERT_OK(Evaluate(op, input_tensor, output_tensor));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_data));
+}
+
+TYPED_TEST(QuantizedUnaryElementWiseTest, QuantizedPerAxisWithAbs) {
+  using StorageT = typename TypeParam::StorageT;
+  using ExpressedT = typename TypeParam::ExpressedT;
+
+  const Shape shape({4, 3, 2});
+  const int quantized_dimension = 2;
+  const size_t rank = shape.Rank();
+  const Axis quantized_dimension_size = shape.Dim(quantized_dimension);
+  const size_t quantization_stride = [&] {
+    size_t res = 1;
+    for (int64_t i = rank - 1; i > quantized_dimension; --i) {
+      res *= shape.Dim(i);
+    }
+    return res;
+  }();
+  Vector<StorageT> input_data = IotaBuffer<TypeParam::kStorage>(shape);
+  Vector<StorageT> output_data(shape.NumElements());
+  Vector<StorageT> zero_points_data = RandomBuffer<TypeParam::kStorage>(
+      /*shape=*/Shape({shape.Dim(2)}), /*min=*/static_cast<StorageT>(-5),
+      /*max=*/static_cast<StorageT>(5));
+  Vector<ExpressedT> scales_data = RandomBuffer<TypeParam::kExpressed>(
+      /*shape=*/Shape({shape.Dim(2)}), /*min=*/static_cast<ExpressedT>(1),
+      /*max=*/static_cast<ExpressedT>(3));
+  const QuantizedElementTypePerAxis tensor_type = QuantizedElementTypePerAxis(
+      TypeParam::kStorage, zero_points_data, TypeParam::kExpressed, scales_data,
+      quantized_dimension);
+  Tensor input_tensor{
+      .type = QuantizedPerAxisTensorType{.shape = shape,
+                                         .element_type = tensor_type},
+      .data = input_data.data()};
+  Tensor output_tensor{
+      .type = QuantizedPerAxisTensorType{.shape = shape,
+                                         .element_type = tensor_type},
+      .data = output_data.data()};
+
+  Vector<StorageT> expected_data(shape.NumElements());
+  absl::c_transform(
+      input_data, expected_data.begin(),
+      [&, element_index = 0ull, quantization_index = 0ull](auto v) mutable {
+        const StorageT zero_point = zero_points_data[quantization_index];
+        const ExpressedT scale = scales_data[quantization_index];
+
+        if (++element_index >= quantization_stride) {
+          element_index = 0;
+          if (++quantization_index >= quantized_dimension_size) {
+            quantization_index = 0;
+          }
+        }
+        const ExpressedT dequantized_input = Dequantize(v, zero_point, scale);
+        const ExpressedT dequantized_res = Abs()(dequantized_input);
+        return Quantize<TypeParam::kStorage, TypeParam::kExpressed>(
+            dequantized_res, zero_point, ExpressedT(1) / scale);
+      });
+
+  auto op = Create(UnaryElementwiseOp<Abs>::Attributes{}, Abs());
+  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/unary_elementwise_test_util.h

@@ -0,0 +1,262 @@
+/* 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_UNARY_ELEMENTWISE_TEST_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UNARY_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>
+using BaselineMismatchSignedIntegerTypes = ::testing::Types<
+    std::tuple<Op, TestParam<DataType::kSI4>, TestParam<DataType::kSI8>>,
+    std::tuple<Op, TestParam<DataType::kSI4>, TestParam<DataType::kSI16>>,
+    std::tuple<Op, TestParam<DataType::kSI8>, TestParam<DataType::kSI4>>,
+    std::tuple<Op, TestParam<DataType::kSI8>, TestParam<DataType::kSI16>>,
+    std::tuple<Op, TestParam<DataType::kSI16>, TestParam<DataType::kSI4>>,
+    std::tuple<Op, TestParam<DataType::kSI16>, TestParam<DataType::kSI8>>>;
+
+// Lists couples of unmatched baseline element types.
+template <class Op>
+using UnaryElementwiseConstraint1Types = ::testing::Types<
+    std::tuple<Op, TestParam<DataType::kF16>, TestParam<DataType::kBF16>>,
+    std::tuple<Op, TestParam<DataType::kF16>, TestParam<DataType::kF32>>,
+    std::tuple<Op, TestParam<DataType::kBF16>, TestParam<DataType::kF16>>,
+    std::tuple<Op, TestParam<DataType::kBF16>, TestParam<DataType::kF32>>,
+    std::tuple<Op, TestParam<DataType::kF32>, TestParam<DataType::kF16>>,
+    std::tuple<Op, TestParam<DataType::kF32>, TestParam<DataType::kBF16>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI4, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF16>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI8, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI16, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI4, DataType::kF32>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kBF16>>>,
+    std::tuple<Op, PerTensor<TestParam<DataType::kSI16, DataType::kF32>>,
+               PerTensor<TestParam<DataType::kSI8, DataType::kF32>>>>;
+
+// Tests that the input shape is compared to the output shape and that it is
+// propagated if needed.
+
+template <class Op>
+class UnaryElementwiseOpShapePropagationTest : public ::testing::Test {
+ protected:
+  void SetOutputShape(Shape shape) {
+    output_tensor_.shape() = std::move(shape);
+  }
+  bool InputAndOutputShapesAreEqual() const {
+    return input_tensor_.shape() == output_tensor_.shape();
+  }
+
+  Op op_ = Create(typename Op::Attributes{});
+  Tensor input_tensor_ = {
+      .type = TensorType{.shape = Shape({2, 3, 4}),
+                         .element_type = SupportedOpDataType<Op>::kStorageType},
+      .data = nullptr};
+  Tensor output_tensor_ = {
+      .type = TensorType{.shape = Shape(),
+                         .element_type = SupportedOpDataType<Op>::kStorageType},
+      .data = nullptr};
+};
+
+TYPED_TEST_SUITE_P(UnaryElementwiseOpShapePropagationTest);
+
+TYPED_TEST_P(UnaryElementwiseOpShapePropagationTest, ShapePropagationWorks) {
+  ASSERT_TRUE(this->output_tensor_.shape().empty());
+  EXPECT_OK(Prepare(this->op_, this->input_tensor_, this->output_tensor_));
+  EXPECT_THAT(this->output_tensor_.shape(),
+              ::testing::ElementsAreArray(this->input_tensor_.shape()));
+}
+
+TYPED_TEST_P(UnaryElementwiseOpShapePropagationTest,
+             SmallerOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3}));
+  ASSERT_FALSE(this->InputAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->input_tensor_, this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shape."));
+}
+
+TYPED_TEST_P(UnaryElementwiseOpShapePropagationTest,
+             BiggerOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3, 4, 5}));
+  ASSERT_FALSE(this->InputAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->input_tensor_, this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shape."));
+}
+
+TYPED_TEST_P(UnaryElementwiseOpShapePropagationTest,
+             IncompatibleOutputShapeRaisesAnError) {
+  this->SetOutputShape(Shape({2, 3, 5}));
+  ASSERT_FALSE(this->InputAndOutputShapesAreEqual());
+  EXPECT_EQ(
+      Prepare(this->op_, this->input_tensor_, this->output_tensor_),
+      absl::FailedPreconditionError("The specified output tensor shape is not "
+                                    "compatible with the input shape."));
+}
+
+REGISTER_TYPED_TEST_SUITE_P(UnaryElementwiseOpShapePropagationTest,
+                            ShapePropagationWorks,
+                            SmallerOutputShapeRaisesAnError,
+                            BiggerOutputShapeRaisesAnError,
+                            IncompatibleOutputShapeRaisesAnError);
+
+// Tests that the baseline element type of the input and output tensors is the
+// same.
+template <class T>
+class UnaryElementwiseSameBaselineElementTypeConstraintTest
+    : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(UnaryElementwiseSameBaselineElementTypeConstraintTest);
+
+TYPED_TEST_P(UnaryElementwiseSameBaselineElementTypeConstraintTest,
+             DifferentInputOutputStorageTypesRaiseAnError) {
+  using Op = std::tuple_element_t<0, TypeParam>;
+  using OperandTypeDesc = std::tuple_element_t<1, TypeParam>;
+  using ResultTypeDesc = std::tuple_element_t<2, TypeParam>;
+  const Shape shape({2, 3, 4});
+  Tensor input_tensor{.type = TensorTypeFor(OperandTypeDesc{}, shape),
+                      .data = nullptr};
+  Tensor output_tensor{.type = TensorTypeFor(ResultTypeDesc{}, shape),
+                       .data = nullptr};
+  auto op = Create(typename Op::Attributes{});
+  const absl::Status status = Prepare(op, input_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(
+    UnaryElementwiseSameBaselineElementTypeConstraintTest,
+    DifferentInputOutputStorageTypesRaiseAnError);
+
+// Tests that unsupported types are detected during when `Prepare` is called.
+template <class T>
+class UnaryElementwiseUnsupportedTypeTest : public ::testing::Test {};
+
+TYPED_TEST_SUITE_P(UnaryElementwiseUnsupportedTypeTest);
+
+TYPED_TEST_P(UnaryElementwiseUnsupportedTypeTest, 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, 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(UnaryElementwiseUnsupportedTypeTest,
+                            PrepareRaisesAnError);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UNARY_ELEMENTWISE_TEST_UTIL_H_

tensorflow/lite/experimental/shlo/ops/util.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 "tensorflow/lite/experimental/shlo/ops/util.h"
+
+#include <string>
+#include <variant>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/data_type.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+absl::Status Propagate(const Shape& input_shape, Shape& output_shape) {
+  if (output_shape.Dimensions().empty()) {
+    output_shape = input_shape;
+  } else if (output_shape != input_shape) {
+    return absl::FailedPreconditionError(
+        "The specified output tensor shape is not compatible with the input "
+        "shape.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status Propagate(const Shape& lhs_shape, const Shape& rhs_shape,
+                       Shape& output_shape) {
+  if (lhs_shape != rhs_shape) {
+    return absl::FailedPreconditionError(
+        "The LHS and RHS shapes are incompatible.");
+  } else if (output_shape.Dimensions().empty()) {
+    output_shape = lhs_shape;
+  } else if (output_shape != lhs_shape) {
+    return absl::FailedPreconditionError(
+        "The specified output tensor shape is not compatible with the input "
+        "shapes.");
+  }
+  return absl::OkStatus();
+}
+
+bool IsBoolTensor(const Tensor& tensor) {
+  return !tensor.IsQuantized() && IsBool(tensor.StorageType());
+}
+
+bool IsSignedIntTensor(const Tensor& tensor) {
+  return !tensor.IsQuantized() && IsSignedInteger(tensor.StorageType());
+}
+
+bool IsUnsignedIntTensor(const Tensor& tensor) {
+  return !tensor.IsQuantized() && IsUnsignedInteger(tensor.StorageType());
+}
+
+bool IsIntTensor(const Tensor& tensor) {
+  return !tensor.IsQuantized() && IsInteger(tensor.StorageType());
+}
+
+bool IsFloatTensor(const Tensor& tensor) {
+  return !tensor.IsQuantized() && IsFloat(tensor.StorageType());
+}
+
+bool IsQuantizedPerTensorTensor(const Tensor& tensor) {
+  return tensor.IsPerTensorQuantized();
+}
+
+bool IsQuantizedPerAxisTensor(const Tensor& tensor) {
+  return tensor.IsPerAxisQuantized();
+}
+
+absl::Status CheckSameBaselineType(CheckCtx ctx, const Tensor& tensor1,
+                                   const Tensor& tensor2) {
+  if (BaselineType(tensor1.element_type()) !=
+      BaselineType(tensor2.element_type())) {
+    std::string tensor1_type_repr =
+        std::visit([](auto v) -> std::string { return ToString(v); },
+                   tensor1.element_type());
+    std::string tensor2_type_repr =
+        std::visit([](auto v) -> std::string { return ToString(v); },
+                   tensor2.element_type());
+    return absl::FailedPreconditionError(
+        "stablehlo." + ctx.op_name +
+        ": baseline type constraint is not satisfied " + tensor1_type_repr +
+        " and " + tensor2_type_repr + ".");
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace shlo_ref

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

@@ -0,0 +1,101 @@
+/* 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_UTIL_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UTIL_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+#define SHLO_REF_RETURN_ON_ERROR(EXPR)    \
+  if (absl::Status s = (EXPR); !s.ok()) { \
+    return s;                             \
+  }
+
+// Propagates the input shape to the output shape.
+//
+// If the output shape is already populated, checks that is it compatible with
+// the input.
+absl::Status Propagate(const Shape& input_shape, Shape& output_shape);
+
+// Propagates the input shapes to the output shape.
+//
+// If the output shape is already populated, checks that is it compatible with
+// the inputs.
+absl::Status Propagate(const Shape& lhs_shape, const Shape& rhs_shape,
+                       Shape& output_shape);
+
+// Provides context information for the `Check*` functions error messages.
+struct CheckCtx {
+  explicit CheckCtx(std::string name) : op_name(name) {}
+  // The operation that requested the check.
+  std::string op_name;
+};
+
+// Checks that the `tensor` element type is supported by one the the `checks`
+// functions.
+//
+// Returns a failed precondition error when no check succeeds.
+//
+// The check functions should have the following signature.
+//
+// ```
+// bool Check(const Tensor& tensor);
+// ```
+template <class... CheckFuncs>
+absl::Status CheckSupportedTypes(CheckCtx ctx, const Tensor& tensor,
+                                 CheckFuncs&&... checks) {
+  if ((static_cast<CheckFuncs&&>(checks)(tensor) || ...)) {
+    return absl::OkStatus();
+  }
+  std::string tensor_type_repr = std::visit(
+      [](auto v) -> std::string { return ToString(v); }, tensor.element_type());
+  return absl::FailedPreconditionError("stablehlo." + ctx.op_name +
+                                       ": Unsupported tensor type (" +
+                                       tensor_type_repr + ").");
+}
+
+// Returns true if the tensor's storage type is boolean.
+bool IsBoolTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is a signed integer type.
+bool IsSignedIntTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is an unsigned integer type.
+bool IsUnsignedIntTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is an integer type.
+bool IsIntTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is an floating point type.
+bool IsFloatTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is quantized per tensor.
+bool IsQuantizedPerTensorTensor(const Tensor& tensor);
+
+// Returns true if the tensor's storage type is quantized per axis.
+bool IsQuantizedPerAxisTensor(const Tensor& tensor);
+
+// Checks that both tensors have the same baseline element type.
+absl::Status CheckSameBaselineType(CheckCtx ctx, const Tensor& tensor1,
+                                   const Tensor& tensor2);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_UTIL_H_

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

@@ -0,0 +1,50 @@
+/* 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/util.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/shape.h"
+#include "tensorflow/lite/experimental/shlo/status_matcher.h"
+
+using testing::ElementsAreArray;
+
+namespace shlo_ref {
+namespace {
+
+TEST(UtilTest, PropagateToEmptyShape) {
+  const Shape input_shape({2, 3, 4, 5});
+  Shape output_shape;
+
+  EXPECT_OK(Propagate(input_shape, output_shape));
+  EXPECT_THAT(output_shape.Dimensions(),
+              ElementsAreArray(input_shape.Dimensions()));
+}
+
+TEST(UtilTest, PropagateToIncompatibleShapeFails) {
+  const Shape input_shape({2, 3, 4, 5});
+  Shape output_shape({2, 3, 4, 6});
+
+  EXPECT_THAT(
+      Propagate(input_shape, output_shape),
+      absl::FailedPreconditionError(
+          "The specified output tensor shape is not compatible with the input "
+          "shape."));
+  EXPECT_THAT(output_shape.Dimensions(), ElementsAreArray({2, 3, 4, 6}));
+}
+
+}  // namespace
+}  // namespace shlo_ref

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

@@ -0,0 +1,68 @@
+/* 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 xor
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/ops/xor.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 Xor : std::bit_xor<void> {};
+
+template <>
+struct Xor<DataType::kI1> {
+  template <class T>
+  bool operator()(T lhs, T rhs) const {
+    return static_cast<bool>(lhs) != static_cast<bool>(rhs);
+  }
+};
+
+XorOp Create(XorOp::Attributes) { return {}; }
+
+absl::Status Prepare(XorOp& 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("xor"), lhs, IsBoolTensor, IsIntTensor));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("xor"), lhs, output));
+  SHLO_REF_RETURN_ON_ERROR(CheckSameBaselineType(CheckCtx("xor"), rhs, output));
+  return absl::OkStatus();
+}
+
+absl::Status Evaluate(XorOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output) {
+  if (IsIntTensor(lhs)) {
+    // Note: all the integer types share the same implementation.
+    Xor<DataType::kSI32> xor_func;
+    DISPATCH_INT(detail::EvaluateNoQuantization, lhs.tensor_element_type(),
+                 xor_func, lhs, rhs, output);
+  } else if (IsBoolTensor(lhs)) {
+    Xor<DataType::kI1> xor_func;
+    detail::EvaluateNoQuantization<DataType::kI1>(xor_func, lhs, rhs, output);
+    return absl::OkStatus();
+  }
+  return absl::FailedPreconditionError(
+      "stablehlo.xor: Unsupported tensor type in Evaluate.");
+}
+
+}  // namespace shlo_ref

tensorflow/lite/experimental/shlo/ops/xor.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 xor
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_XOR_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_XOR_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/lite/experimental/shlo/tensor.h"
+
+namespace shlo_ref {
+
+struct XorOp {
+  struct Attributes {};
+};
+
+XorOp Create(XorOp::Attributes);
+absl::Status Prepare(XorOp& op, const Tensor& lhs, const Tensor& rhs,
+                     Tensor& output);
+absl::Status Evaluate(XorOp& op, const Tensor& lhs, const Tensor& rhs,
+                      Tensor& output);
+
+}  // namespace shlo_ref
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_SHLO_OPS_XOR_H_

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

@@ -0,0 +1,112 @@
+/* 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 xor
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/experimental/shlo/ops/xor.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<XorOp> {
+  static std::string Get() { return "Xor"; }
+};
+
+template <DataType>
+struct Xor : std::bit_xor<void> {};
+
+template <>
+struct Xor<DataType::kI1> {
+  template <class T>
+  bool operator()(T lhs, T rhs) const {
+    return static_cast<bool>(lhs) != static_cast<bool>(rhs);
+  }
+};
+
+template <>
+struct SupportedOpDataType<XorOp> {
+  static constexpr DataType kStorageType = DataType::kSI32;
+};
+
+namespace {
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Xor, BinaryElementwiseOpShapePropagationTest,
+                               XorOp, TestParamNames);
+
+using MultipyBaselineContraintTypes = BinaryElementwiseBaselineConstraintTypes<
+    XorOp, ConcatTypes<BoolTestType, BaselineConstraintIntTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(
+    Xor, BinaryElementwiseSameBaselineElementTypeConstraintTest,
+    MultipyBaselineContraintTypes, TestParamNames);
+
+using UnsupportedTypes =
+    WithOpTypes<XorOp, ConcatTypes<FloatTestTypes, PerTensorQuantizedTestTypes,
+                                   PerAxisQuantizedTestTypes>>;
+
+INSTANTIATE_TYPED_TEST_SUITE_P(Xor, BinaryElementwiseUnsupportedTypeTest,
+                               UnsupportedTypes, TestParamNames);
+
+using SupportedTypes = ConcatTypes<BoolTestType, IntTestTypes>;
+
+template <class T>
+struct XorTest : ::testing::Test {};
+
+TYPED_TEST_SUITE(XorTest, SupportedTypes, TestParamNames);
+
+TYPED_TEST(XorTest, 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(),
+                    Xor<TypeParam::kStorage>());
+
+  auto op = Create(XorOp::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