From bfccee3de4bd22f3ad23aad113793c597fd18769 Mon Sep 17 00:00:00 2001
From: Scott Wolchok <swolchok@meta.com>
Date: Thu, 23 Jan 2025 10:47:41 -0800
Subject: [PATCH] Support Half/BFloat16 in sign (#7866)

Partial fix for #7748.
---
 kernels/portable/cpu/op_sign.cpp |  2 +-
 kernels/test/op_sign_test.cpp    | 29 +++++++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/kernels/portable/cpu/op_sign.cpp b/kernels/portable/cpu/op_sign.cpp
index af3225d477..a003811461 100644
--- a/kernels/portable/cpu/op_sign.cpp
+++ b/kernels/portable/cpu/op_sign.cpp
@@ -39,7 +39,7 @@ Tensor& sign_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
   if (in.scalar_type() == exec_aten::ScalarType::Bool) {
     memcpy(out.mutable_data_ptr(), in.const_data_ptr(), in.nbytes());
   } else {
-    ET_SWITCH_REAL_TYPES(in.scalar_type(), ctx, "sign.out", CTYPE, [&] {
+    ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, "sign.out", CTYPE, [&] {
       apply_unary_map_fn(
           [](const CTYPE val_in) {
             if (std::isnan(val_in)) {
diff --git a/kernels/test/op_sign_test.cpp b/kernels/test/op_sign_test.cpp
index e411675b19..8e7fd8b453 100644
--- a/kernels/test/op_sign_test.cpp
+++ b/kernels/test/op_sign_test.cpp
@@ -25,22 +25,31 @@ class OpSignTest : public OperatorTest {
   Tensor& op_sign_out(const Tensor& self, Tensor& out) {
     return torch::executor::aten::sign_outf(context_, self, out);
   }
+
+  template <typename CTYPE, ScalarType DTYPE>
+  void test_et_dtype() {
+    TensorFactory<DTYPE> tf;
+
+    const auto infinity = std::numeric_limits<CTYPE>::infinity();
+    const auto nan = std::numeric_limits<CTYPE>::quiet_NaN();
+    Tensor in = tf.make({1, 7}, {-infinity, -3., -1.5, 0., 1.5, nan, infinity});
+    Tensor out = tf.zeros({1, 7});
+    Tensor expected = tf.make({1, 7}, {-1., -1., -1., 0., 1., nan, 1.});
+
+    Tensor ret = op_sign_out(in, out);
+
+    EXPECT_TENSOR_EQ(out, ret);
+    EXPECT_TENSOR_CLOSE(out, expected);
+  }
 };
 
 TEST_F(OpSignTest, ETSanityCheckFloat) {
   if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
     GTEST_SKIP() << "ATen returns 0 on NAN input";
   }
-  TensorFactory<ScalarType::Float> tf;
-
-  Tensor in = tf.make({1, 7}, {-INFINITY, -3., -1.5, 0., 1.5, NAN, INFINITY});
-  Tensor out = tf.zeros({1, 7});
-  Tensor expected = tf.make({1, 7}, {-1., -1., -1., 0., 1., NAN, 1.});
-
-  Tensor ret = op_sign_out(in, out);
-
-  EXPECT_TENSOR_EQ(out, ret);
-  EXPECT_TENSOR_CLOSE(out, expected);
+#define TEST_ENTRY(ctype, dtype) test_et_dtype<ctype, ScalarType::dtype>();
+  ET_FORALL_FLOATHBF16_TYPES(TEST_ENTRY);
+#undef TEST_ENTRY
 }
 
 TEST_F(OpSignTest, ATenSanityCheckFloat) {