diff --git a/tensorflow/core/kernels/mlir_generated/BUILD b/tensorflow/core/kernels/mlir_generated/BUILD
index 0862b37848a..bff0811c63e 100644
--- a/tensorflow/core/kernels/mlir_generated/BUILD
+++ b/tensorflow/core/kernels/mlir_generated/BUILD
@@ -197,6 +197,7 @@ tf_cuda_cc_test(
         "no_cuda_asan",  # TODO(b/171341759): re-enable.
     ],
     deps = [
+        ":gpu_ops_test_util",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:tensorflow",
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
index 30c34c1405c..1f2bc105e1a 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
+++ b/tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h
@@ -48,10 +48,12 @@ absl::InlinedVector<T, 10> RepeatInputToMatchShape(
   return result;
 }
 
-/// Helper functions to get default input values.
+/// Helper functions to get default input shapes.
 
 TensorShape DefaultInputShape();
 
+/// Helper functions to get default input data.
+
 template <typename T,
           std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
                            bool> = true>
@@ -72,17 +74,10 @@ T DefaultScalarInput() {
   return static_cast<T>(true);
 }
 
-template <typename T>
-absl::InlinedVector<T, 10> InfZeroInput() {
-  return InputAsVector<T, double>({-std::numeric_limits<double>::infinity(),
-                                   -0.1, -0.0, 0.0, 0.1,
-                                   std::numeric_limits<float>::infinity()});
-}
-
 template <typename T,
           std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
                            bool> = true>
-absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
+absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name = "") {
   // Only generate values less than the bitwidth of the data type.
   if (op_name == "LeftShift" || op_name == "RightShift") {
     auto max_shift = sizeof(T) * 8 - 1;
@@ -96,16 +91,65 @@ absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
 template <typename T, std::enable_if_t<
                           llvm::is_one_of<T, Eigen::half, float, double>::value,
                           bool> = true>
-absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
+absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name = "") {
   return InputAsVector<T, double>({-18.0, -9.0, -1e-6, -0.0, 0.0, 1e-6, 0.1,
                                    0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
 }
 
 template <typename T,
           std::enable_if_t<llvm::is_one_of<T, bool>::value, bool> = true>
-absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name) {
+absl::InlinedVector<T, 10> DefaultInput(absl::string_view op_name = "") {
   return InputAsVector<T, bool>({true, false, true, true, false});
 }
+
+/// Helper functions to get more specific input data.
+
+template <typename T, std::enable_if_t<
+                          llvm::is_one_of<T, Eigen::half, float, double>::value,
+                          bool> = true>
+absl::InlinedVector<std::complex<T>, 10> DefaultComplexInput() {
+  auto input = test::DefaultInput<T>();
+  absl::InlinedVector<std::complex<T>, 10> complex_input;
+  for (T value : input) {
+    complex_input.emplace_back(value, -value);
+  }
+  return complex_input;
+}
+
+template <typename T, std::enable_if_t<
+                          llvm::is_one_of<T, Eigen::half, float, double>::value,
+                          bool> = true>
+absl::InlinedVector<T, 10> NearZeroAndExtremeInput() {
+  return InputAsVector<T, double>({-std::numeric_limits<double>::infinity(),
+                                   -0.1, -0.0, 0.0, 0.1,
+                                   std::numeric_limits<float>::infinity()});
+}
+
+template <typename T,
+          std::enable_if_t<llvm::is_one_of<T, int8, int16, int32, int64>::value,
+                           bool> = true>
+absl::InlinedVector<T, 10> NearZeroAndExtremeInput() {
+  return InputAsVector<T, T>({std::numeric_limits<T>::min(),
+                              std::numeric_limits<T>::min() + 1, -1, 0, 1,
+                              std::numeric_limits<T>::max()});
+}
+
+template <typename T, std::enable_if_t<
+                          llvm::is_one_of<T, Eigen::half, float, double>::value,
+                          bool> = true>
+absl::InlinedVector<T, 10> DefaultInputGreaterThanZero() {
+  return test::InputAsVector<T, double>({18.0, 9.0, 1e-6, 1.0, 0.1, 1e-6, 0.1,
+                                         0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
+}
+
+template <typename T, std::enable_if_t<
+                          llvm::is_one_of<T, Eigen::half, float, double>::value,
+                          bool> = true>
+absl::InlinedVector<T, 10> DefaultInputGreaterOrEqualToZero() {
+  return test::InputAsVector<T, double>({18.0, 9.0, 1e-6, 0.0, 0.1, 1e-6, 0.1,
+                                         0.2, 0.3, 0.5, 0.7, 0.9, 9.0, 18.0});
+}
+
 }  // namespace test
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
index 64c548e2e41..a3308b5c01f 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/kernels/mlir_generated/gpu_ops_test_util.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/test.h"
@@ -44,20 +45,10 @@ class GpuUnaryOpTest : public OpsTestBase {
     SetDevice(tensorflow::DEVICE_GPU, std::move(device_gpu));
   }
 
-  // 'T' is the input type, 'RT' is the input type for the callback function,
-  // 'OutT' is the output type, 'ROutT' is the output type for the callback
-  // function. In most cases it is enough to just provide the input type,
-  // because all the types are the same.
-  template <typename T, typename RT = T, typename OutT = T, typename ROutT = RT>
-  void Run(std::vector<int64> input_shape, absl::InlinedVector<T, 10> input,
-           const std::string op_name, ROutT (*expected_callback)(RT),
-           bool expect_equal = true, bool add_tout = false,
-           bool expect_buffer_reuse = true, bool add_t = true) {
-    assert(std::accumulate(input_shape.begin(), input_shape.end(), 1,
-                           std::multiplies<int64>()) == input.size() &&
-           "Expected input length to equal to shape's number of elements.");
-
-    TensorShape shape(input_shape);
+  template <typename T, typename OutT>
+  void SetOpKernel(const std::string& op_name, const TensorShape& shape,
+                   const absl::InlinedVector<T, 10>& input, bool add_t,
+                   bool add_tout) {
     NodeDefBuilder builder("some_name", op_name);
     builder.Input(FakeInput(DataTypeToEnum<T>::v()));
     if (add_t) {
@@ -70,6 +61,15 @@ class GpuUnaryOpTest : public OpsTestBase {
 
     TF_ASSERT_OK(InitOp());
     AddInputFromArray<T>(shape, input);
+  }
+
+  template <typename T, typename OutT>
+  void RunAndExpectResult(const std::string& op_name, const TensorShape& shape,
+                          const absl::InlinedVector<T, 10>& input,
+                          const absl::InlinedVector<OutT, 10>& expected_output,
+                          bool add_t, bool add_tout, bool expect_buffer_reuse,
+                          bool expect_equal) {
+    SetOpKernel<T, OutT>(op_name, shape, input, add_t, add_tout);
     TF_ASSERT_OK(RunOpKernel());
 
     // Assert buffer reuse if expected.
@@ -81,13 +81,7 @@ class GpuUnaryOpTest : public OpsTestBase {
 
     // Assert expected results.
     Tensor expected_tensor(allocator(), DataTypeToEnum<OutT>::value, shape);
-    absl::InlinedVector<OutT, 14> expected;
-    expected.reserve(input.size());
-    for (const T& inp : input) {
-      expected.push_back(
-          static_cast<OutT>(expected_callback(static_cast<RT>(inp))));
-    }
-    test::FillValues<OutT>(&expected_tensor, expected);
+    test::FillValues<OutT>(&expected_tensor, expected_output);
     if (expect_equal) {
       test::ExpectEqual(expected_tensor, *GetOutput(0));
     } else {
@@ -95,241 +89,225 @@ class GpuUnaryOpTest : public OpsTestBase {
     }
   }
 
-  // Some helper functions to get default input values.
+  template <typename T, typename BaselineT, typename OutT,
+            typename BaselineOutT>
+  void Test(const std::string op_name, const TensorShape& shape,
+            absl::InlinedVector<T, 10> input,
+            BaselineOutT (*baseline_callback)(BaselineT),
+            bool expect_equal = true, bool add_tout = false,
+            bool expect_buffer_reuse = true, bool add_t = true) {
+    // Prepare inputs and compute expected results.
+    auto repeated_input =
+        test::RepeatInputToMatchShape(input, shape.num_elements());
+    absl::InlinedVector<OutT, 10> expected_output =
+        ComputeExpectedOutput<T, BaselineT, OutT, BaselineOutT>(
+            repeated_input, baseline_callback);
 
-  std::vector<int64> DefaultInputShape() { return std::vector<int64>{2, 7}; }
-
-  template <typename T>
-  absl::InlinedVector<T, 10> DefaultInput() {
-    return InputAsVector<T>({-18.0, -9.0, -1e-6, -0.0, 0.0, 1e-6, 0.1, 0.2, 0.3,
-                             0.5, 0.7, 0.9, 9.0, 18.0});
-  }
-
-  template <typename T>
-  absl::InlinedVector<std::complex<T>, 10> DefaultComplexInput() {
-    auto input = DefaultInput<T>();
-    absl::InlinedVector<std::complex<T>, 10> complex_input;
-    for (T value : input) {
-      complex_input.emplace_back(value, -value);
-    }
-    return complex_input;
-  }
-
-  template <typename T>
-  absl::InlinedVector<T, 10> DefaultInputGreaterThanZero() {
-    return InputAsVector<T>({18.0, 9.0, 1e-6, 1.0, 0.1, 1e-6, 0.1, 0.2, 0.3,
-                             0.5, 0.7, 0.9, 9.0, 18.0});
-  }
-
-  template <typename T>
-  absl::InlinedVector<T, 10> DefaultInputGreaterOrEqualToZero() {
-    return InputAsVector<T>({18.0, 9.0, 1e-6, 0.0, 0.1, 1e-6, 0.1, 0.2, 0.3,
-                             0.5, 0.7, 0.9, 9.0, 18.0});
+    RunAndExpectResult<T, OutT>(op_name, shape, repeated_input, expected_output,
+                                add_t, add_tout, expect_buffer_reuse,
+                                expect_equal);
   }
 
  private:
-  template <typename T>
-  absl::InlinedVector<T, 10> InputAsVector(
-      std::initializer_list<double> input) {
-    absl::InlinedVector<T, 10> result;
-    result.reserve(input.size());
-    for (const auto& value : input) {
-      result.push_back(static_cast<T>(value));
+  template <typename T, typename BaselineT, typename OutT,
+            typename BaselineOutT>
+  absl::InlinedVector<OutT, 10> ComputeExpectedOutput(
+      absl::InlinedVector<T, 10> input,
+      BaselineOutT (*baseline_callback)(BaselineT)) {
+    absl::InlinedVector<OutT, 10> expected_output;
+    for (int i = 0; i < input.size(); i++) {
+      auto arg = static_cast<BaselineT>(input[i]);
+      auto result = static_cast<OutT>(baseline_callback(arg));
+      expected_output.push_back(result);
     }
-    return result;
+    return expected_output;
   }
 };
 
 /// Test `tf.Abs`.
 
 TEST_F(GpuUnaryOpTest, AbsFloat) {
-  Run<float>(
-      /*input_shape=*/{2, 3},
-      /*input=*/
-      {-std::numeric_limits<float>::infinity(), -0.1f, -0.0f, 0.0f, 0.1f,
-       std::numeric_limits<float>::infinity()},
-      /*op_name=*/"Abs",
-      /*expected_callback=*/std::abs,
+  Test<float, float, float, float>(
+      /*op_name=*/"Abs", test::DefaultInputShape(),
+      test::NearZeroAndExtremeInput<float>(),
+      /*baseline_callback=*/std::abs,
       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, AbsDouble) {
-  Run<double>(
-      /*input_shape=*/{2, 3},
-      /*input=*/
-      {-std::numeric_limits<double>::infinity(), -0.1, -0.0, 0.0, 0.1,
-       std::numeric_limits<double>::infinity()},
-      /*op_name=*/"Abs",
-      /*expected_callback=*/std::abs,
+  Test<double, double, double, double>(
+      /*op_name=*/"Abs", test::DefaultInputShape(),
+      test::NearZeroAndExtremeInput<double>(),
+      /*baseline_callback=*/std::abs,
       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, AbsHalf) {
-  Run<Eigen::half, float>(
-      /*input_shape=*/{2, 3},
-      /*input=*/
-      {static_cast<Eigen::half>(-std::numeric_limits<double>::infinity()),
-       static_cast<Eigen::half>(-0.1), static_cast<Eigen::half>(-0.0),
-       static_cast<Eigen::half>(0.0), static_cast<Eigen::half>(0.1),
-       static_cast<Eigen::half>(std::numeric_limits<double>::infinity())},
-      /*op_name=*/"Abs",
-      /*expected_callback=*/std::abs,
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Abs", test::DefaultInputShape(),
+      test::NearZeroAndExtremeInput<Eigen::half>(),
+      /*baseline_callback=*/std::abs,
       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, AbsInt32) {
-  Run<int32>(
-      /*input_shape=*/{2, 3},
-      /*input=*/
-      {std::numeric_limits<int32>::min(), std::numeric_limits<int32>::min() + 1,
-       -1, 0, 1, std::numeric_limits<int32>::max()},
-      /*op_name=*/"Abs",
-      /*expected_callback=*/std::abs,
+  Test<int32, int32, int32, int32>(
+      /*op_name=*/"Abs", test::DefaultInputShape(),
+      test::NearZeroAndExtremeInput<int32>(),
+      /*baseline_callback=*/std::abs,
       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, AbsInt64) {
-  Run<int64>(
-      /*input_shape=*/{2, 3},
-      /*input=*/
-      {std::numeric_limits<int64>::min(), std::numeric_limits<int64>::min() + 1,
-       -1, 0, 1, std::numeric_limits<int64>::max()},
-      /*op_name=*/"Abs",
-      /*expected_callback=*/std::abs,
+  Test<int64, int64, int64, int64>(
+      /*op_name=*/"Abs", test::DefaultInputShape(),
+      test::NearZeroAndExtremeInput<int64>(),
+      /*baseline_callback=*/std::abs,
       /*expect_equal=*/true);
 }
 
 /// Test `tf.Ceil`.
 
 TEST_F(GpuUnaryOpTest, CeilFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Ceil",
-             /*expected_callback=*/std::ceil,
-             /*expect_equal=*/true);
+  Test<float, float, float, float>(
+      /*op_name=*/"Ceil", test::DefaultInputShape(),
+      test::DefaultInput<float>("Ceil"),
+      /*baseline_callback=*/std::ceil,
+      /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, CeilDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Ceil",
-              /*expected_callback=*/std::ceil,
-              /*expect_equal=*/true);
+  Test<double, double, double, double>(
+      /*op_name=*/"Ceil", test::DefaultInputShape(),
+      test::DefaultInput<double>(),
+      /*baseline_callback=*/std::ceil,
+      /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, CeilHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Ceil",
-                          /*expected_callback=*/std::ceil,
-                          /*expect_equal=*/true);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Ceil", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::ceil,
+      /*expect_equal=*/true);
 }
 
 /// Test `tf.Conj`.
 
 TEST_F(GpuUnaryOpTest, ConjFloat) {
-  Run<std::complex<float>, const std::complex<float>&, std::complex<float>,
-      std::complex<float>>(DefaultInputShape(), DefaultComplexInput<float>(),
-                           /*op_name=*/"Conj",
-                           /*expected_callback=*/std::conj,
-                           /*expect_equal=*/false,
-                           /*add_tout=*/false,
-                           /*expect_buffer_reuse=*/false);
-}
-
-TEST_F(GpuUnaryOpTest, ConjDouble) {
-  Run<std::complex<double>, const std::complex<double>&, std::complex<double>,
-      std::complex<double>>(DefaultInputShape(), DefaultComplexInput<double>(),
-                            /*op_name=*/"Conj",
-                            /*expected_callback=*/std::conj,
+  Test<std::complex<float>, const std::complex<float>&, std::complex<float>,
+       std::complex<float>>(/*op_name=*/"Conj", test::DefaultInputShape(),
+                            test::DefaultComplexInput<float>(),
+                            /*baseline_callback=*/std::conj,
                             /*expect_equal=*/false,
                             /*add_tout=*/false,
                             /*expect_buffer_reuse=*/false);
 }
 
+TEST_F(GpuUnaryOpTest, ConjDouble) {
+  Test<std::complex<double>, const std::complex<double>&, std::complex<double>,
+       std::complex<double>>(
+      /*op_name=*/"Conj", test::DefaultInputShape(),
+      test::DefaultComplexInput<double>(),
+      /*baseline_callback=*/std::conj,
+      /*expect_equal=*/false,
+      /*add_tout=*/false,
+      /*expect_buffer_reuse=*/false);
+}
+
 /// Test `tf.Cos`.
 
 TEST_F(GpuUnaryOpTest, CosFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Cos",
-             /*expected_callback=*/std::cos,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(
+      /*op_name=*/"Cos", test::DefaultInputShape(), test::DefaultInput<float>(),
+      /*baseline_callback=*/std::cos,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, CosDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Cos",
-              /*expected_callback=*/std::cos,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(/*op_name=*/"Cos",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/std::cos,
+                                       /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, CosHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Cos",
-                          /*expected_callback=*/std::cos,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Cos", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::cos,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.Exp`.
 
 TEST_F(GpuUnaryOpTest, ExpFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Exp",
-             /*expected_callback=*/std::exp,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(/*op_name=*/"Exp", test::DefaultInputShape(),
+                                   test::DefaultInput<float>(),
+                                   /*baseline_callback=*/std::exp,
+                                   /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, ExpDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Exp",
-              /*expected_callback=*/std::exp,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(/*op_name=*/"Exp",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/std::exp,
+                                       /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, ExpHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Exp",
-                          /*expected_callback=*/std::exp,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Exp", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::exp,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.Floor`.
 
 TEST_F(GpuUnaryOpTest, FloorFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Floor",
-             /*expected_callback=*/std::floor,
-             /*expect_equal=*/true);
+  Test<float, float, float, float>(/*op_name=*/"Floor",
+                                   test::DefaultInputShape(),
+                                   test::DefaultInput<float>(),
+                                   /*baseline_callback=*/std::floor,
+                                   /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, FloorDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Floor",
-              /*expected_callback=*/std::floor,
-              /*expect_equal=*/true);
+  Test<double, double, double, double>(/*op_name=*/"Floor",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/std::floor,
+                                       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, FloorHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Floor",
-                          /*expected_callback=*/std::floor,
-                          /*expect_equal=*/true);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Floor", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::floor,
+      /*expect_equal=*/true);
 }
 
 /// Test `tf.Imag`.
 
 TEST_F(GpuUnaryOpTest, ImagFloat) {
-  Run<std::complex<float>, const std::complex<float>&, float, float>(
-      DefaultInputShape(), DefaultComplexInput<float>(),
-      /*op_name=*/"Imag",
-      /*expected_callback=*/std::imag,
+  Test<std::complex<float>, const std::complex<float>&, float, float>(
+      /*op_name=*/"Imag", test::DefaultInputShape(),
+      test::DefaultComplexInput<float>(),
+      /*baseline_callback=*/std::imag,
       /*expect_equal=*/false,
       /*add_tout=*/true,
       /*expect_buffer_reuse=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, ImagDouble) {
-  Run<std::complex<double>, const std::complex<double>&, double, double>(
-      DefaultInputShape(), DefaultComplexInput<double>(),
-      /*op_name=*/"Imag",
-      /*expected_callback=*/std::imag,
+  Test<std::complex<double>, const std::complex<double>&, double, double>(
+      /*op_name=*/"Imag", test::DefaultInputShape(),
+      test::DefaultComplexInput<double>(),
+      /*baseline_callback=*/std::imag,
       /*expect_equal=*/false,
       /*add_tout=*/true,
       /*expect_buffer_reuse=*/false);
@@ -338,64 +316,65 @@ TEST_F(GpuUnaryOpTest, ImagDouble) {
 /// Test `tf.IsInf`.
 
 // TODO(b/162575339): The tests currently still fails with CUDA_ILLEGAL_ADDRESS
-// when run with unranked kernels.
+// when Test with unranked kernels.
 TEST_F(GpuUnaryOpTest, DISABLED_IsInfFloat) {
-  Run<float, float, bool, bool>(DefaultInputShape(), DefaultInput<float>(),
-                                /*op_name=*/"IsInf",
-                                /*expected_callback=*/std::isinf,
-                                /*expect_equal=*/true);
+  Test<float, float, bool, bool>(/*op_name=*/"IsInf", test::DefaultInputShape(),
+                                 test::DefaultInput<float>(),
+                                 /*baseline_callback=*/std::isinf,
+                                 /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, DISABLED_IsInfDouble) {
   // Workaround for gcc bug, it would fail with "unresolved overloaded function
   // type" if passing std::isinf with type double. So we use type float for
   // comparing expected values.
-  Run<double, float, bool, bool>(DefaultInputShape(), DefaultInput<double>(),
-                                 /*op_name=*/"IsInf",
-                                 /*expected_callback=*/std::isinf,
-                                 /*expect_equal=*/true);
+  Test<double, float, bool, bool>(/*op_name=*/"IsInf",
+                                  test::DefaultInputShape(),
+                                  test::DefaultInput<double>(),
+                                  /*baseline_callback=*/std::isinf,
+                                  /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, DISABLED_IsInfHalf) {
-  Run<Eigen::half, float, bool, bool>(DefaultInputShape(),
-                                      DefaultInput<Eigen::half>(),
-                                      /*op_name=*/"IsInf",
-                                      /*expected_callback=*/std::isinf,
-                                      /*expect_equal=*/true);
+  Test<Eigen::half, float, bool, bool>(/*op_name=*/"IsInf",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<Eigen::half>(),
+                                       /*baseline_callback=*/std::isinf,
+                                       /*expect_equal=*/true);
 }
 
 /// Test `tf.Log`.
 
 TEST_F(GpuUnaryOpTest, LogFloat) {
-  Run<float>(DefaultInputShape(), DefaultInputGreaterThanZero<float>(),
-             /*op_name=*/"Log",
-             /*expected_callback=*/std::log,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(/*op_name=*/"Log", test::DefaultInputShape(),
+                                   test::DefaultInputGreaterThanZero<float>(),
+                                   /*baseline_callback=*/std::log,
+                                   /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, LogDouble) {
-  Run<double>(DefaultInputShape(), DefaultInputGreaterThanZero<double>(),
-              /*op_name=*/"Log",
-              /*expected_callback=*/std::log,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(
+      /*op_name=*/"Log", test::DefaultInputShape(),
+      test::DefaultInputGreaterThanZero<double>(),
+      /*baseline_callback=*/std::log,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, LogHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(),
-                          /*input=*/
-                          DefaultInputGreaterThanZero<Eigen::half>(),
-                          /*op_name=*/"Log",
-                          /*expected_callback=*/std::log,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Log", test::DefaultInputShape(),
+      test::DefaultInputGreaterThanZero<Eigen::half>(),
+      /*baseline_callback=*/std::log,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.LogicalNot`
 
 TEST_F(GpuUnaryOpTest, LogicalNot) {
-  Run<bool, bool, bool, bool>(
-      DefaultInputShape(), DefaultInput<bool>(),
-      /*op_name=*/"LogicalNot",
-      /*expected_callback=*/[](bool v) { return !v; },
+  Test<bool, bool, bool, bool>(
+      /*op_name=*/"LogicalNot", test::DefaultInputShape(),
+      test::DefaultInput<bool>(),
+      /*baseline_callback=*/[](bool v) { return !v; },
       /*expect_equal=*/true,
       /*add_tout=*/false,
       /*expect_buffer_reuse=*/true,
@@ -406,69 +385,71 @@ TEST_F(GpuUnaryOpTest, LogicalNot) {
 
 /// Reference implementation.
 template <typename T>
-T expected_neg(T x) {
+T baseline_neg(T x) {
   return -x;
 }
 
 TEST_F(GpuUnaryOpTest, NegFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Neg",
-             /*expected_callback=*/expected_neg,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(
+      /*op_name=*/"Neg", test::DefaultInputShape(), test::DefaultInput<float>(),
+      /*baseline_callback=*/baseline_neg,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, NegDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Neg",
-              /*expected_callback=*/expected_neg,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(
+      /*op_name=*/"Neg", test::DefaultInputShape(),
+      test::DefaultInput<double>(),
+      /*baseline_callback=*/baseline_neg,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, NegHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Neg",
-                          /*expected_callback=*/expected_neg,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Neg", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/baseline_neg,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, NegInt8) {
-  Run<int8>(DefaultInputShape(), DefaultInput<int8>(),
-            /*op_name=*/"Neg",
-            /*expected_callback=*/expected_neg,
-            /*expect_equal=*/true);
+  Test<int8, int8, int8, int8>(
+      /*op_name=*/"Neg", test::DefaultInputShape(), test::DefaultInput<int8>(),
+      /*baseline_callback=*/baseline_neg,
+      /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, NegInt16) {
-  Run<int16>(DefaultInputShape(), DefaultInput<int16>(),
-             /*op_name=*/"Neg",
-             /*expected_callback=*/expected_neg,
-             /*expect_equal=*/true);
+  Test<int16, int16, int16, int16>(/*op_name=*/"Neg", test::DefaultInputShape(),
+                                   test::DefaultInput<int16>(),
+                                   /*baseline_callback=*/baseline_neg,
+                                   /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, NegInt64) {
-  Run<int64>(DefaultInputShape(), DefaultInput<int64>(),
-             /*op_name=*/"Neg",
-             /*expected_callback=*/expected_neg,
-             /*expect_equal=*/true);
+  Test<int64, int64, int64, int64>(/*op_name=*/"Neg", test::DefaultInputShape(),
+                                   test::DefaultInput<int64>(),
+                                   /*baseline_callback=*/baseline_neg,
+                                   /*expect_equal=*/true);
 }
 
 /// Test `tf.Real`.
 
 TEST_F(GpuUnaryOpTest, RealFloat) {
-  Run<std::complex<float>, const std::complex<float>&, float, float>(
-      DefaultInputShape(), DefaultComplexInput<float>(),
-      /*op_name=*/"Real",
-      /*expected_callback=*/std::real,
+  Test<std::complex<float>, const std::complex<float>&, float, float>(
+      /*op_name=*/"Real", test::DefaultInputShape(),
+      test::DefaultComplexInput<float>(),
+      /*baseline_callback=*/std::real,
       /*expect_equal=*/false,
       /*add_tout=*/true,
       /*expect_buffer_reuse=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, RealDouble) {
-  Run<std::complex<double>, const std::complex<double>&, double, double>(
-      DefaultInputShape(), DefaultComplexInput<double>(),
-      /*op_name=*/"Real",
-      /*expected_callback=*/std::real,
+  Test<std::complex<double>, const std::complex<double>&, double, double>(
+      /*op_name=*/"Real", test::DefaultInputShape(),
+      test::DefaultComplexInput<double>(),
+      /*baseline_callback=*/std::real,
       /*expect_equal=*/false,
       /*add_tout=*/true,
       /*expect_buffer_reuse=*/false);
@@ -478,141 +459,153 @@ TEST_F(GpuUnaryOpTest, RealDouble) {
 
 /// Reference implementation.
 template <typename T>
-T expected_rsqrt(T x) {
+T baseline_rsqrt(T x) {
   return 1.0 / std::sqrt(x);
 }
 
 TEST_F(GpuUnaryOpTest, RsqrtFloat) {
-  Run<float>(DefaultInputShape(), DefaultInputGreaterThanZero<float>(),
-             /*op_name=*/"Rsqrt",
-             /*expected_callback=*/expected_rsqrt,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(/*op_name=*/"Rsqrt",
+                                   test::DefaultInputShape(),
+                                   test::DefaultInputGreaterThanZero<float>(),
+                                   /*baseline_callback=*/baseline_rsqrt,
+                                   /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, RsqrtDouble) {
-  Run<double>(DefaultInputShape(), DefaultInputGreaterThanZero<double>(),
-              /*op_name=*/"Rsqrt",
-              /*expected_callback=*/expected_rsqrt,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(
+      /*op_name=*/"Rsqrt", test::DefaultInputShape(),
+      test::DefaultInputGreaterThanZero<double>(),
+      /*baseline_callback=*/baseline_rsqrt,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, RsqrtHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(),
-                          /*input=*/
-                          DefaultInputGreaterThanZero<Eigen::half>(),
-                          /*op_name=*/"Rsqrt",
-                          /*expected_callback=*/expected_rsqrt,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Rsqrt", test::DefaultInputShape(),
+      test::DefaultInputGreaterThanZero<Eigen::half>(),
+      /*baseline_callback=*/baseline_rsqrt,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.Sign`.
 
 // Reference implementation
 template <typename T>
-T expected_sign(T x) {
+T baseline_sign(T x) {
   if (x == 0) return 0;
   if (x < 0) return -1;
   return 1;
 }
 
 TEST_F(GpuUnaryOpTest, SignFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Sign",
-             /*expected_callback=*/expected_sign,
-             /*expect_equal=*/true);
+  Test<float, float, float, float>(/*op_name=*/"Sign",
+                                   test::DefaultInputShape(),
+                                   test::DefaultInput<float>(),
+                                   /*baseline_callback=*/baseline_sign,
+                                   /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, SignDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Sign",
-              /*expected_callback=*/expected_sign,
-              /*expect_equal=*/true);
+  Test<double, double, double, double>(/*op_name=*/"Sign",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/baseline_sign,
+                                       /*expect_equal=*/true);
 }
 
 TEST_F(GpuUnaryOpTest, SignHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Sign",
-                          /*expected_callback=*/expected_sign,
-                          // TODO(b/162577610): We should actually use true
-                          // here. This requires returning 0.0 for input -0.0.
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Sign", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*expected_callback=*/baseline_sign,
+      // TODO(b/162577610): We should actually use true
+      // here. This requires returning 0.0 for input -0.0.
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, SignInt64) {
-  Run<int64>(DefaultInputShape(), DefaultInput<int64>(),
-             /*op_name=*/"Sign",
-             /*expected_callback=*/expected_sign,
-             /*expect_equal=*/true);
+  Test<int64, int64, int64, int64>(
+      /*op_name=*/"Sign", test::DefaultInputShape(),
+      test::DefaultInput<int64>(),
+      /*expected_callback=*/baseline_sign,
+      /*expect_equal=*/true);
 }
 
 /// Test `tf.Sin`.
 
 TEST_F(GpuUnaryOpTest, SinFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Sin",
-             /*expected_callback=*/std::sin,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(/*op_name=*/"Sin", test::DefaultInputShape(),
+                                   test::DefaultInput<float>(),
+                                   /*baseline_callback=*/std::sin,
+                                   /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, SinDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Sin",
-              /*expected_callback=*/std::sin,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(/*op_name=*/"Sin",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/std::sin,
+                                       /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, SinHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Sin",
-                          /*expected_callback=*/std::sin,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Sin", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::sin,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.Sqrt`.
 
 TEST_F(GpuUnaryOpTest, SqrtFloat) {
-  Run<float>(DefaultInputShape(), DefaultInputGreaterOrEqualToZero<float>(),
-             /*op_name=*/"Sqrt",
-             /*expected_callback=*/std::sqrt,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(
+      /*op_name=*/"Sqrt", test::DefaultInputShape(),
+      test::DefaultInputGreaterOrEqualToZero<float>(),
+      /*baseline_callback=*/std::sqrt,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, SqrtDouble) {
-  Run<double>(DefaultInputShape(), DefaultInputGreaterOrEqualToZero<double>(),
-              /*op_name=*/"Sqrt",
-              /*expected_callback=*/std::sqrt,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(
+      /*op_name=*/"Sqrt", test::DefaultInputShape(),
+      test::DefaultInputGreaterOrEqualToZero<double>(),
+      /*baseline_callback=*/std::sqrt,
+      /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, SqrtHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(),
-                          DefaultInputGreaterOrEqualToZero<Eigen::half>(),
-                          /*op_name=*/"Sqrt",
-                          /*expected_callback=*/std::sqrt,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Sqrt", test::DefaultInputShape(),
+      test::DefaultInputGreaterOrEqualToZero<Eigen::half>(),
+      /*baseline_callback=*/std::sqrt,
+      /*expect_equal=*/false);
 }
 
 /// Test `tf.Tanh`.
 
 TEST_F(GpuUnaryOpTest, TanhFloat) {
-  Run<float>(DefaultInputShape(), DefaultInput<float>(),
-             /*op_name=*/"Tanh",
-             /*expected_callback=*/std::tanh,
-             /*expect_equal=*/false);
+  Test<float, float, float, float>(/*op_name=*/"Tanh",
+                                   test::DefaultInputShape(),
+                                   test::DefaultInput<float>(),
+                                   /*baseline_callback=*/std::tanh,
+                                   /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, TanhDouble) {
-  Run<double>(DefaultInputShape(), DefaultInput<double>(),
-              /*op_name=*/"Tanh",
-              /*expected_callback=*/std::tanh,
-              /*expect_equal=*/false);
+  Test<double, double, double, double>(/*op_name=*/"Tanh",
+                                       test::DefaultInputShape(),
+                                       test::DefaultInput<double>(),
+                                       /*baseline_callback=*/std::tanh,
+                                       /*expect_equal=*/false);
 }
 
 TEST_F(GpuUnaryOpTest, TanhHalf) {
-  Run<Eigen::half, float>(DefaultInputShape(), DefaultInput<Eigen::half>(),
-                          /*op_name=*/"Tanh",
-                          /*expected_callback=*/std::tanh,
-                          /*expect_equal=*/false);
+  Test<Eigen::half, float, Eigen::half, float>(
+      /*op_name=*/"Tanh", test::DefaultInputShape(),
+      test::DefaultInput<Eigen::half>(),
+      /*baseline_callback=*/std::tanh,
+      /*expect_equal=*/false);
 }
 
 }  // namespace