Move ti.zeros() to dpctl_ext/tensor and reuse it in dpnp

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -50,6 +50,7 @@
     ones_like,
     tril,
     triu,
+    zeros,
 )
 from dpctl_ext.tensor._indexing_functions import (
     extract,
@@ -105,4 +106,5 @@
     "tril",
     "triu",
     "where",
+    "zeros",
 ]

dpctl_ext/tensor/_ctors.py

@@ -1091,7 +1091,7 @@ def eye(
     n_cols = n_rows if n_cols is None else operator.index(n_cols)
     k = operator.index(k)
     if k >= n_cols or -k >= n_rows:
-        return dpt.zeros(
+        return dpt_ext.zeros(
             (n_rows, n_cols),
             dtype=dtype,
             order=order,
@@ -1720,7 +1720,7 @@ def tril(x, /, *, k=0):
         )
         _manager.add_event_pair(hev, cpy_ev)
     elif k < -shape[nd - 2]:
-        res = dpt.zeros(
+        res = dpt_ext.zeros(
             x.shape,
             dtype=x.dtype,
             order=order,
@@ -1784,7 +1784,7 @@ def triu(x, /, *, k=0):
 
     q = x.sycl_queue
     if k > shape[nd - 1]:
-        res = dpt.zeros(
+        res = dpt_ext.zeros(
             x.shape,
             dtype=x.dtype,
             order=order,
@@ -1821,3 +1821,72 @@ def triu(x, /, *, k=0):
         _manager.add_event_pair(hev, triu_ev)
 
     return res
+
+
+def zeros(
+    shape,
+    *,
+    dtype=None,
+    order="C",
+    device=None,
+    usm_type="device",
+    sycl_queue=None,
+):
+    """
+    Returns a new :class:`dpctl.tensor.usm_ndarray` having a specified
+    shape and filled with zeros.
+
+    Args:
+        shape (Tuple[int], int):
+            Dimensions of the array to be created.
+        dtype (optional):
+            data type of the array. Can be typestring,
+            a :class:`numpy.dtype` object, :mod:`numpy` char string,
+            or a NumPy scalar type. Default: ``None``
+        order ("C", or "F"):
+            memory layout for the array. Default: ``"C"``
+        device (optional): array API concept of device where the output array
+            is created. ``device`` can be ``None``, a oneAPI filter selector
+            string, an instance of :class:`dpctl.SyclDevice` corresponding to
+            a non-partitioned SYCL device, an instance of
+            :class:`dpctl.SyclQueue`, or a :class:`dpctl.tensor.Device` object
+            returned by :attr:`dpctl.tensor.usm_ndarray.device`.
+            Default: ``None``
+        usm_type (``"device"``, ``"shared"``, ``"host"``, optional):
+            The type of SYCL USM allocation for the output array.
+            Default: ``"device"``
+        sycl_queue (:class:`dpctl.SyclQueue`, optional):
+            The SYCL queue to use
+            for output array allocation and copying. ``sycl_queue`` and
+            ``device`` are complementary arguments, i.e. use one or another.
+            If both are specified, a :exc:`TypeError` is raised unless both
+            imply the same underlying SYCL queue to be used. If both are
+            ``None``, a cached queue targeting default-selected device is
+            used for allocation and population. Default: ``None``
+
+    Returns:
+        usm_ndarray:
+            Constructed array initialized with zeros.
+    """
+    if not isinstance(order, str) or len(order) == 0 or order[0] not in "CcFf":
+        raise ValueError(
+            "Unrecognized order keyword value, expecting 'F' or 'C'."
+        )
+    order = order[0].upper()
+    dpctl.utils.validate_usm_type(usm_type, allow_none=False)
+    sycl_queue = normalize_queue_device(sycl_queue=sycl_queue, device=device)
+    dtype = _get_dtype(dtype, sycl_queue)
+    _ensure_native_dtype_device_support(dtype, sycl_queue.sycl_device)
+    res = dpt.usm_ndarray(
+        shape,
+        dtype=dtype,
+        buffer=usm_type,
+        order=order,
+        buffer_ctor_kwargs={"queue": sycl_queue},
+    )
+    _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
+    # populating new allocation, no dependent events
+    hev, zeros_ev = ti._zeros_usm_ndarray(res, sycl_queue)
+    _manager.add_event_pair(hev, zeros_ev)
+
+    return res

dpnp/dpnp_container.py

@@ -35,12 +35,11 @@
 
 """
 
-import dpctl.tensor as dpt
 import dpctl.utils as dpu
 
 # TODO: revert to `import dpctl.tensor...`
 # when dpnp fully migrates dpctl/tensor
-import dpctl_ext.tensor as dpt_ext
+import dpctl_ext.tensor as dpt
 import dpnp
 from dpnp.dpnp_array import dpnp_array
 
@@ -75,7 +74,7 @@ def arange(
         sycl_queue=sycl_queue, device=device
     )
 
-    array_obj = dpt_ext.arange(
+    array_obj = dpt.arange(
         start,
         stop=stop,
         step=step,
@@ -103,7 +102,7 @@ def asarray(
 
     """Converts incoming 'x1' object to 'dpnp_array'."""
     if isinstance(x1, (list, tuple, range)):
-        array_obj = dpt_ext.asarray(
+        array_obj = dpt.asarray(
             x1,
             dtype=dtype,
             copy=copy,
@@ -122,7 +121,7 @@ def asarray(
             x1_obj, device=device, sycl_queue=sycl_queue
         )
 
-        array_obj = dpt_ext.asarray(
+        array_obj = dpt.asarray(
             x1_obj,
             dtype=dtype,
             copy=copy,
@@ -143,7 +142,7 @@ def copy(x1, /, *, order="K"):
     if order is None:
         order = "K"
 
-    array_obj = dpt_ext.copy(dpnp.get_usm_ndarray(x1), order=order)
+    array_obj = dpt.copy(dpnp.get_usm_ndarray(x1), order=order)
     return dpnp_array._create_from_usm_ndarray(array_obj)
 
 
@@ -165,7 +164,7 @@ def empty(
         order = "C"
 
     """Creates `dpnp_array` from uninitialized USM allocation."""
-    array_obj = dpt_ext.empty(
+    array_obj = dpt.empty(
         shape,
         dtype=dtype,
         order=order,
@@ -196,7 +195,7 @@ def eye(
         order = "C"
 
     """Creates `dpnp_array` with ones on the `k`th diagonal."""
-    array_obj = dpt_ext.eye(
+    array_obj = dpt.eye(
         N,
         M,
         k=k,
@@ -231,7 +230,7 @@ def full(
         fill_value = fill_value.get_array()
 
     """Creates `dpnp_array` having a specified shape, filled with fill_value."""
-    array_obj = dpt_ext.full(
+    array_obj = dpt.full(
         shape,
         fill_value,
         dtype=dtype,
@@ -260,7 +259,7 @@ def ones(
         order = "C"
 
     """Creates `dpnp_array` of ones with the given shape, dtype, and order."""
-    array_obj = dpt_ext.ones(
+    array_obj = dpt.ones(
         shape,
         dtype=dtype,
         order=order,
@@ -272,13 +271,13 @@ def ones(
 
 def tril(x1, /, *, k=0):
     """Creates `dpnp_array` as lower triangular part of an input array."""
-    array_obj = dpt_ext.tril(dpnp.get_usm_ndarray(x1), k=k)
+    array_obj = dpt.tril(dpnp.get_usm_ndarray(x1), k=k)
     return dpnp_array._create_from_usm_ndarray(array_obj)
 
 
 def triu(x1, /, *, k=0):
     """Creates `dpnp_array` as upper triangular part of an input array."""
-    array_obj = dpt_ext.triu(dpnp.get_usm_ndarray(x1), k=k)
+    array_obj = dpt.triu(dpnp.get_usm_ndarray(x1), k=k)
     return dpnp_array._create_from_usm_ndarray(array_obj)