fix wrong return types for some of the math dialect stmts (#474)

Author: kaihsin

src/kirin/dialects/math/__init__.py

@@ -89,15 +89,15 @@ def gamma(x: float) -> float: ...
 
 
 @lowering.wraps(stmts.isfinite)
-def isfinite(x: float) -> float: ...
+def isfinite(x: float) -> bool: ...
 
 
 @lowering.wraps(stmts.isinf)
-def isinf(x: float) -> float: ...
+def isinf(x: float) -> bool: ...
 
 
 @lowering.wraps(stmts.isnan)
-def isnan(x: float) -> float: ...
+def isnan(x: float) -> bool: ...
 
 
 @lowering.wraps(stmts.lgamma)

src/kirin/dialects/math/_gen.py

@@ -56,6 +56,10 @@ def builtin_math_functions():
                 for arg in sig.parameters.keys()
             ]
         )
+        if "is" in name:
+            ret_type = "types.Bool"
+        else:
+            ret_type = "types.Float"
         f.write(
             textwrap.dedent(
                 f"""
@@ -66,7 +70,7 @@ class {name}(ir.Statement):
     name = "{name}"
     traits = frozenset({{ir.Pure(), lowering2.FromPythonCall()}})
 {fields}
-    result: ir.ResultValue = info.result(types.Float)
+    result: ir.ResultValue = info.result({ret_type})
 """
             )
         )
@@ -116,11 +120,15 @@ class MathMethodTable(MethodTable):
     f.write("from kirin import lowering2\n")
 
     for name, obj, sig in builtin_math_functions():
+        if "is" in name:
+            ret_type = "bool"
+        else:
+            ret_type = "float"
         f.write(
             textwrap.dedent(
                 f"""
         @lowering2.wraps(stmts.{name})
-        def {name}({", ".join(f"{arg}: float" for arg in sig.parameters.keys())}) -> float: ...
+        def {name}({", ".join(f"{arg}: {ret_type}" for arg in sig.parameters.keys())}) -> {ret_type}: ...
         """
             )
         )

src/kirin/dialects/math/stmts.py

@@ -204,7 +204,7 @@ class isfinite(ir.Statement):
     name = "isfinite"
     traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     x: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(types.Float)
+    result: ir.ResultValue = info.result(types.Bool)
 
 
 @statement(dialect=dialect)
@@ -214,7 +214,7 @@ class isinf(ir.Statement):
     name = "isinf"
     traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     x: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(types.Float)
+    result: ir.ResultValue = info.result(types.Bool)
 
 
 @statement(dialect=dialect)
@@ -224,7 +224,7 @@ class isnan(ir.Statement):
     name = "isnan"
     traits = frozenset({ir.Pure(), lowering.FromPythonCall()})
     x: ir.SSAValue = info.argument(types.Float)
-    result: ir.ResultValue = info.result(types.Float)
+    result: ir.ResultValue = info.result(types.Bool)
 
 
 @statement(dialect=dialect)