Merge pull request #1571 from IntelPython/improve-coverage-usmarray

Improve coverage usmarray

Author: oleksandr-pavlyk

dpctl/tests/test_usm_ndarray_ctor.py

@@ -462,12 +462,16 @@ def test_ctor_buffer_kwarg():
     with pytest.raises(ValueError):
         dpt.usm_ndarray(10, buffer="invalid_param")
     Xusm = dpt.usm_ndarray((10, 5), dtype="c8")
+    Xusm[...] = 1
     X2 = dpt.usm_ndarray(Xusm.shape, buffer=Xusm, dtype=Xusm.dtype)
-    assert np.array_equal(
-        Xusm.usm_data.copy_to_host(), X2.usm_data.copy_to_host()
-    )
+    Horig_copy = Xusm.usm_data.copy_to_host()
+    H2_copy = X2.usm_data.copy_to_host()
+    assert np.array_equal(Horig_copy, H2_copy)
     with pytest.raises(ValueError):
         dpt.usm_ndarray(10, dtype="i4", buffer=dict())
+    # use device-specific default fp data type
+    X3 = dpt.usm_ndarray(Xusm.shape, buffer=Xusm)
+    assert np.array_equal(Horig_copy, X3.usm_data.copy_to_host())
 
 
 def test_usm_ndarray_props():

dpctl/tests/test_usm_ndarray_indexing.py

@@ -164,7 +164,7 @@ def test_basic_slice10():
 
 
 def _all_equal(it1, it2):
-    return all(dpt.asnumpy(x) == dpt.asnumpy(y) for x, y in zip(it1, it2))
+    return all(bool(x == y) for x, y in zip(it1, it2))
 
 
 def test_advanced_slice1():
@@ -175,8 +175,6 @@ def test_advanced_slice1():
     assert isinstance(y, dpt.usm_ndarray)
     assert y.shape == ii.shape
     assert y.strides == (1,)
-    # FIXME, once usm_ndarray.__equal__ is implemented,
-    # use of asnumpy should be removed
     assert _all_equal(
         (x[ii[k]] for k in range(ii.shape[0])),
         (y[k] for k in range(ii.shape[0])),
@@ -185,8 +183,6 @@ def test_advanced_slice1():
     assert isinstance(y, dpt.usm_ndarray)
     assert y.shape == ii.shape
     assert y.strides == (1,)
-    # FIXME, once usm_ndarray.__equal__ is implemented,
-    # use of asnumpy should be removed
     assert _all_equal(
         (x[ii[k]] for k in range(ii.shape[0])),
         (y[k] for k in range(ii.shape[0])),
@@ -201,8 +197,6 @@ def test_advanced_slice1_negative_strides():
     assert isinstance(y, dpt.usm_ndarray)
     assert y.shape == ii.shape
     assert y.strides == (1,)
-    # FIXME, once usm_ndarray.__equal__ is implemented,
-    # use of asnumpy should be removed
     assert _all_equal(
         (x[ii[k]] for k in range(ii.shape[0])),
         (y[k] for k in range(ii.shape[0])),
@@ -400,6 +394,16 @@ def test_advanced_slice13():
     assert (dpt.asnumpy(y) == dpt.asnumpy(expected)).all()
 
 
+def test_boolean_indexing_validation():
+    get_queue_or_skip()
+    x = dpt.zeros(10, dtype="i4")
+    ii = dpt.ones((2, 5), dtype="?")
+    with pytest.raises(IndexError):
+        x[ii]
+    with pytest.raises(IndexError):
+        x[ii[0, :]]
+
+
 def test_integer_indexing_1d():
     get_queue_or_skip()
     x = dpt.arange(10, dtype="i4")
@@ -482,6 +486,32 @@ def test_TrueFalse_indexing():
         assert y3._pointer == x._pointer
 
 
+def test_mixed_index_getitem():
+    get_queue_or_skip()
+    x = dpt.reshape(dpt.arange(1000, dtype="i4"), (10, 10, 10))
+    i1b = dpt.ones(10, dtype="?")
+    info = x.__array_namespace__().__array_namespace_info__()
+    ind_dt = info.default_dtypes(x.device)["indexing"]
+    i0 = dpt.asarray([0, 2, 3], dtype=ind_dt)[:, dpt.newaxis]
+    i2 = dpt.asarray([3, 4, 7], dtype=ind_dt)[:, dpt.newaxis]
+    y = x[i0, i1b, i2]
+    assert y.shape == (3, dpt.sum(i1b, dtype="i8"))
+
+
+def test_mixed_index_setitem():
+    get_queue_or_skip()
+    x = dpt.reshape(dpt.arange(1000, dtype="i4"), (10, 10, 10))
+    i1b = dpt.ones(10, dtype="?")
+    info = x.__array_namespace__().__array_namespace_info__()
+    ind_dt = info.default_dtypes(x.device)["indexing"]
+    i0 = dpt.asarray([0, 2, 3], dtype=ind_dt)[:, dpt.newaxis]
+    i2 = dpt.asarray([3, 4, 7], dtype=ind_dt)[:, dpt.newaxis]
+    v_shape = (3, int(dpt.sum(i1b, dtype="i8")))
+    canary = 7
+    x[i0, i1b, i2] = dpt.full(v_shape, canary, dtype=x.dtype)
+    assert x[0, 0, 3] == canary
+
+
 @pytest.mark.parametrize(
     "data_dt",
     _all_dtypes,