Fix code generation for numpy scalars in Python target

inducer · Jun 25, 2024 · 2c7f942 · 2c7f942
1 parent 44982ac
commit 2c7f942
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/loopy/target/python.py b/loopy/target/python.py
@@ -25,6 +25,8 @@
 
 from typing import Optional, Sequence, Tuple
 
+import numpy as np
+
 from pymbolic.mapper import Mapper
 from pymbolic.mapper.stringifier import StringifyMapper
 from genpy import Generable, Suite, Collection
@@ -58,7 +60,10 @@ def rec(self, expr, prec, type_context=None, needed_dtype=None):
     __call__ = rec
 
     def map_constant(self, expr, enclosing_prec):
-        return repr(expr)
+        if isinstance(expr, np.generic):
+            return repr(expr).replace("np.", "_lpy_np.")
+        else:
+            return repr(expr)
 
     def map_variable(self, expr, enclosing_prec):
         if expr.name in self.codegen_state.var_subst_map: