Merge pull request #910 from IntelPython/added_tril_and_triu

Added dpctl.tensor.tril and dpctl.tensor.triu feature

Author: oleksandr-pavlyk

dpctl/tensor/__init__.py

@@ -33,6 +33,8 @@
     linspace,
     ones,
     ones_like,
+    tril,
+    triu,
     zeros,
     zeros_like,
 )
@@ -83,4 +85,6 @@
     "to_numpy",
     "asnumpy",
     "from_dlpack",
+    "tril",
+    "triu",
 ]

dpctl/tensor/_ctors.py

@@ -1116,3 +1116,85 @@ def eye(
         hev, _ = ti._eye(k, dst=res, sycl_queue=sycl_queue)
         hev.wait()
     return res
+
+
+def tril(X, k=0):
+    """
+    tril(X: usm_ndarray, k: int) -> usm_ndarray
+
+    Returns the lower triangular part of a matrix (or a stack of matrices) X.
+    """
+    if type(X) is not dpt.usm_ndarray:
+        raise TypeError
+
+    k = operator.index(k)
+
+    # F_CONTIGUOUS = 2
+    order = "F" if (X.flags & 2) else "C"
+
+    shape = X.shape
+    nd = X.ndim
+    if nd < 2:
+        raise ValueError("Array dimensions less than 2.")
+
+    if k >= shape[nd - 1] - 1:
+        res = dpt.empty(
+            X.shape, dtype=X.dtype, order=order, 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
+        )
+    else:
+        res = dpt.empty(
+            X.shape, dtype=X.dtype, order=order, sycl_queue=X.sycl_queue
+        )
+        hev, _ = ti._tril(src=X, dst=res, k=k, sycl_queue=X.sycl_queue)
+        hev.wait()
+
+    return res
+
+
+def triu(X, k=0):
+    """
+    triu(X: usm_ndarray, k: int) -> usm_ndarray
+
+    Returns the upper triangular part of a matrix (or a stack of matrices) X.
+    """
+    if type(X) is not dpt.usm_ndarray:
+        raise TypeError
+
+    k = operator.index(k)
+
+    # F_CONTIGUOUS = 2
+    order = "F" if (X.flags & 2) else "C"
+
+    shape = X.shape
+    nd = X.ndim
+    if nd < 2:
+        raise ValueError("Array dimensions less than 2.")
+
+    if k > shape[nd - 1]:
+        res = dpt.zeros(
+            X.shape, dtype=X.dtype, order=order, 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
+        )
+        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
+        )
+        hev, _ = ti._triu(src=X, dst=res, k=k, sycl_queue=X.sycl_queue)
+        hev.wait()
+
+    return res