(tvm) branch main updated: Fix sqrt/rsqrt Compatibility with Integer Data Types (#17953)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new b77f948489 Fix sqrt/rsqrt Compatibility with Integer Data Types (#17953)
b77f948489 is described below

commit b77f9484898c3f813de661fa058fe2dbef0ee003
Author: Qingchao Shen <qi...@outlook.com>
AuthorDate: Tue May 13 22:32:19 2025 +0800

    Fix sqrt/rsqrt Compatibility with Integer Data Types (#17953)
---
 python/tvm/topi/math.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/tvm/topi/math.py b/python/tvm/topi/math.py
index 865d62683c..5e405cdf6b 100644
--- a/python/tvm/topi/math.py
+++ b/python/tvm/topi/math.py
@@ -515,6 +515,8 @@ def sqrt(x):
     y : tvm.te.Tensor
         The result.
     """
+    if x.dtype.startswith("int"):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
     return te.compute(x.shape, lambda *i: te.sqrt(x(*i)))
 
 
@@ -532,6 +534,8 @@ def rsqrt(x):
     y : tvm.te.Tensor
         The result.
     """
+    if x.dtype.startswith("int"):
+        x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
     return te.compute(x.shape, lambda *i: te.rsqrt(x(*i)))