Tests for new out parameter behavior for add

Author: ndgrigorian

dpctl/tests/elementwise/test_add.py

@@ -381,17 +381,35 @@ def test_add_inplace_dtype_matrix(op1_dtype, op2_dtype):
             dpt.asnumpy(ar1) == np.full(ar1.shape, 2, dtype=ar1.dtype)
         ).all()
 
-        ar3 = dpt.ones(sz, dtype=op1_dtype)
-        ar4 = dpt.ones(2 * sz, dtype=op2_dtype)
-
-        ar3[::-1] += ar4[::2]
+        ar3 = dpt.ones(sz, dtype=op1_dtype)[::-1]
+        ar4 = dpt.ones(2 * sz, dtype=op2_dtype)[::2]
+        ar3 += ar4
         assert (
             dpt.asnumpy(ar3) == np.full(ar3.shape, 2, dtype=ar3.dtype)
         ).all()
-
     else:
         with pytest.raises(TypeError):
             ar1 += ar2
+            dpt.add(ar1, ar2, out=ar1)
+
+    # out is second arg
+    ar1 = dpt.ones(sz, dtype=op1_dtype)
+    ar2 = dpt.ones_like(ar1, dtype=op2_dtype)
+    if _can_cast(ar1.dtype, ar2.dtype, _fp16, _fp64):
+        dpt.add(ar1, ar2, out=ar2)
+        assert (
+            dpt.asnumpy(ar2) == np.full(ar2.shape, 2, dtype=ar2.dtype)
+        ).all()
+
+        ar3 = dpt.ones(sz, dtype=op1_dtype)[::-1]
+        ar4 = dpt.ones(2 * sz, dtype=op2_dtype)[::2]
+        dpt.add(ar3, ar4, out=ar4)
+        assert (
+            dpt.asnumpy(ar4) == np.full(ar4.shape, 2, dtype=ar4.dtype)
+        ).all()
+    else:
+        with pytest.raises(TypeError):
+            dpt.add(ar1, ar2, out=ar2)
 
 
 def test_add_inplace_broadcasting():
@@ -403,6 +421,12 @@ def test_add_inplace_broadcasting():
     m += v
     assert (dpt.asnumpy(m) == np.arange(1, 6, dtype="i4")[np.newaxis, :]).all()
 
+    # check case where second arg is out
+    dpt.add(v, m, out=m)
+    assert (
+        dpt.asnumpy(m) == np.arange(10, dtype="i4")[np.newaxis, 1:10:2]
+    ).all()
+
 
 def test_add_inplace_errors():
     get_queue_or_skip()
@@ -441,7 +465,7 @@ def test_add_inplace_errors():
         ar1 += ar2
 
 
-def test_add_inplace_overlap():
+def test_add_inplace_same_tensors():
     get_queue_or_skip()
 
     ar1 = dpt.ones(10, dtype="i4")
@@ -451,7 +475,13 @@ def test_add_inplace_overlap():
     ar1 = dpt.ones(10, dtype="i4")
     ar2 = dpt.ones(10, dtype="i4")
     dpt.add(ar1, ar2, out=ar1)
+    # all ar1 vals should be 2
     assert (dpt.asnumpy(ar1) == np.full(ar1.shape, 2, dtype="i4")).all()
 
     dpt.add(ar2, ar1, out=ar2)
+    # all ar2 vals should be 3
     assert (dpt.asnumpy(ar2) == np.full(ar2.shape, 3, dtype="i4")).all()
+
+    dpt.add(ar1, ar2, out=ar2)
+    # all ar2 vals should be 5
+    assert (dpt.asnumpy(ar2) == np.full(ar2.shape, 5, dtype="i4")).all()