apache · jikechao · Sep 9, 2023 · Sep 9, 2023 · Sep 9, 2023 · Oct 4, 2023
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
@@ -2116,6 +2116,8 @@ def exp(x):
         The result.
     """
     x = tir.convert(x)
+    if "int" in x.dtype:
+        x = tir.Cast("float32", x)
     return call_intrin(x.dtype, "tir.exp", x)
 
 

diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -85,6 +85,7 @@ def main(x: R.Tensor(("m", "n"), dtype)):
         ("cos", R.cos, topi.cos, "float32"),
         ("cosh", R.cosh, topi.cosh, "float32"),
         ("exp", R.exp, topi.exp, "float32"),
+        ("exp", R.exp, topi.exp, "int32"),
         ("floor", R.floor, topi.floor, "float32"),
         ("floor", R.floor, topi.identity, "int32"),
         ("log", R.log, topi.log, "float32"),