NumPy fusion type coercion fix (#703)

* Fix type coercion in fusion pass

* Add fusion test case

Author: arshajii

codon/cir/transform/numpy/expr.cpp

@@ -63,17 +63,25 @@ types::Type *coerceScalarArray(NumPyType &scalar, NumPyType &array,
     return scalar.getIRBaseType(T);
 }
 
+bool isPythonScalar(NumPyType &t) {
+  if (t.isArray())
+    return false;
+  auto dt = t.dtype;
+  return (dt == NumPyType::NP_TYPE_BOOL || dt == NumPyType::NP_TYPE_I64 ||
+          dt == NumPyType::NP_TYPE_F64 || dt == NumPyType::NP_TYPE_C128);
+}
+
 template <typename E>
 types::Type *decideTypes(E *expr, NumPyType &lhs, NumPyType &rhs,
                          NumPyPrimitiveTypes &T) {
   // Special case(s)
   if (expr->op == E::NP_OP_COPYSIGN)
     return expr->type.getIRBaseType(T);
 
-  if (lhs.isArray() && !rhs.isArray())
+  if (lhs.isArray() && isPythonScalar(rhs))
     return coerceScalarArray(rhs, lhs, T);
 
-  if (!lhs.isArray() && rhs.isArray())
+  if (isPythonScalar(lhs) && rhs.isArray())
     return coerceScalarArray(lhs, rhs, T);
 
   auto *t1 = lhs.getIRBaseType(T);

test/numpy/test_routines.codon

@@ -2479,3 +2479,12 @@ def test_dragon4_interface():
         assert equal(fpos(f('1.001', tp), precision=1, trim='-'), "1")
 
 test_dragon4_interface()
+
+@test
+def test_fusion_coercion():
+    # more extensive tests in test_fusion.codon
+    a = np.array([1.0, 2.0])
+    b = np.multiply(a, np.float32(2)) + np.float32(0.5)
+    assert np.array_equal(b, [2.5, 4.5])
+
+test_fusion_coercion()