Merge pull request #1062 from IntelPython/fix_tril_triu_usm_type

Added missing usm_type to tril() and triu() functions.

Author: oleksandr-pavlyk

dpctl/tensor/_ctors.py

@@ -1247,19 +1247,31 @@ def tril(X, k=0):
 
     if k >= shape[nd - 1] - 1:
         res = dpt.empty(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
         hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
             src=X, dst=res, sycl_queue=X.sycl_queue
         )
         hev.wait()
     elif k < -shape[nd - 2]:
         res = dpt.zeros(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
     else:
         res = dpt.empty(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
         hev, _ = ti._tril(src=X, dst=res, k=k, sycl_queue=X.sycl_queue)
         hev.wait()
@@ -1290,19 +1302,31 @@ def triu(X, k=0):
 
     if k > shape[nd - 1]:
         res = dpt.zeros(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
     elif k <= -shape[nd - 2] + 1:
         res = dpt.empty(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
         hev, _ = ti._copy_usm_ndarray_into_usm_ndarray(
             src=X, dst=res, sycl_queue=X.sycl_queue
         )
         hev.wait()
     else:
         res = dpt.empty(
-            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+            X.shape,
+            dtype=X.dtype,
+            order=order,
+            usm_type=X.usm_type,
+            sycl_queue=X.sycl_queue,
         )
         hev, _ = ti._triu(src=X, dst=res, k=k, sycl_queue=X.sycl_queue)
         hev.wait()

dpctl/tests/test_usm_ndarray_ctor.py

@@ -1503,6 +1503,28 @@ def test_triu(dtype):
     assert np.array_equal(Ynp, dpt.asnumpy(Y))
 
 
+@pytest.mark.parametrize("tri_fn", [dpt.tril, dpt.triu])
+@pytest.mark.parametrize("usm_type", ["device", "shared", "host"])
+def test_tri_usm_type(tri_fn, usm_type):
+    q = get_queue_or_skip()
+    dtype = dpt.uint16
+
+    shape = (2, 3, 4, 5, 5)
+    size = np.prod(shape)
+    X = dpt.reshape(
+        dpt.arange(size, dtype=dtype, usm_type=usm_type, sycl_queue=q), shape
+    )
+    Y = tri_fn(X)  # main execution branch
+    assert Y.usm_type == X.usm_type
+    assert Y.sycl_queue == q
+    Y = tri_fn(X, k=-6)  # special case of Y == X
+    assert Y.usm_type == X.usm_type
+    assert Y.sycl_queue == q
+    Y = tri_fn(X, k=6)  # special case of Y == 0
+    assert Y.usm_type == X.usm_type
+    assert Y.sycl_queue == q
+
+
 def test_tril_slice():
     q = get_queue_or_skip()