pass all arguments to NaN constructor

inducer · Sep 23, 2023 · a9b57f4 · a9b57f4
1 parent da3c0a1
commit a9b57f4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/loopy/target/c/codegen/expression.py b/loopy/target/c/codegen/expression.py
@@ -431,10 +431,8 @@ def map_constant(self, expr, type_context):
                  " The generated code will be equivalent with the added benefit"
                  " of sound pickling/unpickling of kernel objects.")
             from pymbolic.primitives import NaN
-            if not isinstance(expr, np.generic):
-                return self.map_nan(NaN(), type_context)
-            else:
-                return self.map_nan(NaN(expr.dtype.type), type_context)
+            data_type = expr.dtype.type if isinstance(expr, np.generic) else None
+            return self.map_nan(NaN(data_type), type_context)
         elif np.isneginf(expr):
             return -p.Variable("INFINITY")
         elif np.isinf(expr):

diff --git a/test/test_target.py b/test/test_target.py
@@ -400,7 +400,7 @@ def test_nan_support(ctx_factory, target):
          lp.Assignment(parse("c"), parse("isnan(3.14)")),
          lp.Assignment(parse("d"), parse("isnan(0.0)")),
          lp.Assignment(parse("e"), NaN(np.float32)),
-         lp.Assignment(parse("f"), Variable("isnan")(NaN())),
+         lp.Assignment(parse("f"), Variable("isnan")(NaN(None))),
          lp.Assignment(parse("g"), NaN(np.complex64)),
          lp.Assignment(parse("h"), NaN(np.complex128)),
          ],