add more tests

Author: neelay893

test/dialects/vmath/test_desugar.py

@@ -18,11 +18,6 @@ def add_scalar_rhs_typed(x: IList[float, Any], y: float):
     return x + y
 
 
-@basic.union([vmath])
-def add_two_lists():
-    return add_kernel(x=[0, 1, 2], y=[3, 4, 5])
-
-
 @basic.union([vmath])(aggressive=True, typeinfer=True)
 def add_scalar_lhs():
     return add_kernel(x=3.0, y=[3.0, 4, 5])
@@ -41,4 +36,49 @@ def test_add_scalar_lhs():
 def test_typed_kernel_add():
     VMathDesugar(add_scalar_rhs_typed.dialects).unsafe_run(add_scalar_rhs_typed)
     add_scalar_rhs_typed.print()
-    print(add_scalar_rhs_typed(IList([0, 1, 2]), 3.1))
+    res = add_scalar_rhs_typed(IList([0, 1, 2]), 3.1)
+    assert np.allclose(np.asarray(res), np.asarray([3.1, 4.1, 5.1]))
+
+
+@basic.union([vmath])
+def add_two_lists():
+    return add_kernel(x=[0, 1, 2], y=[3, 4, 5])
+
+
+def test_add_lists():
+    VMathDesugar(add_two_lists.dialects).unsafe_run(add_two_lists)
+    res = add_two_lists()
+    assert np.allclose(np.asarray(res), np.array([0, 1, 2, 3, 4, 5]))
+
+
+@basic.union([vmath])
+def sub_scalar_rhs_typed(x: IList[float, Any], y: float):
+    return x - y
+
+
+def test_sub_scalar_typed():
+    VMathDesugar(sub_scalar_rhs_typed.dialects).unsafe_run(sub_scalar_rhs_typed)
+    res = sub_scalar_rhs_typed(IList([0, 1, 2]), 3.1)
+    assert np.allclose(np.asarray(res), np.asarray([-3.1, -2.1, -1.1]))
+
+
+@basic.union([vmath])
+def mult_scalar_lhs_typed(x: float, y: IList[float, Any]):
+    return x * y
+
+
+def test_mult_scalar_typed():
+    VMathDesugar(mult_scalar_lhs_typed.dialects).unsafe_run(mult_scalar_lhs_typed)
+    res = mult_scalar_lhs_typed(3, IList([0, 1, 2]))
+    assert np.allclose(np.asarray(res), np.asarray([0, 3, 6]))
+
+
+@basic.union([vmath])
+def div_scalar_lhs_typed(x: float, y: IList[float, Any]):
+    return x / y
+
+
+def test_div_scalar_typed():
+    VMathDesugar(div_scalar_lhs_typed.dialects).unsafe_run(div_scalar_lhs_typed)
+    res = div_scalar_lhs_typed(3, IList([1, 1.5, 2]))
+    assert np.allclose(np.asarray(res), np.asarray([3, 2, 1.5]))