fix sign return type

test/inductor/test_torchinductor.py

@@ -3111,6 +3111,13 @@ def fn(x):
             (torch.randn([1, 2, 6, 6]),),
         )
 
+    def test_sign_dtype(self):
+        def fn(x):
+            y = torch.sign(x)
+            return torch.tanh(y)
+
+        self.common(fn, (torch.randn([1, 2, 6, 6]),))
+
     def test_fmod(self):
         def fn(a, b):
             return torch.fmod(a, b), torch.fmod(3.0 * a, b) - 2.0

torch/_inductor/codegen/common.py

@@ -260,12 +260,6 @@ def reciprocal(x):
     def square(x):
         return ops.mul(x, x)
 
-    @staticmethod
-    def sign(x):
-        left = ops.where(ops.lt("0", x), "1", "0")
-        right = ops.where(ops.lt(x, "0"), "1", "0")
-        return ops.sub(left, right)
-
     @staticmethod
     def bitwise_not(x):
         return f"~{ExprPrinter.paren(x)}"

torch/_inductor/codegen/triton.py

@@ -417,6 +417,13 @@ def floordiv(a, b):
         rem = f"{a} % {b}"
         return f"tl.where(({a} < 0) != ({b} < 0), tl.where({rem} != 0, {quot} - 1, {quot}), {quot})"
 
+    @staticmethod
+    def sign(x):
+        left = ops.where(ops.lt("0", x), 1, 0)
+        right = ops.where(ops.lt(x, "0"), 1, 0)
+        sub = ops.sub(left, right)
+        return f"{sub}.to({x}.dtype)"
+
     @staticmethod
     def trunc(x):
         return f"tl.math.trunc({x})"