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

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -46,6 +46,7 @@
     full_like,
     linspace,
     meshgrid,
+    ones,
     tril,
     triu,
 )
@@ -88,6 +89,7 @@
     "linspace",
     "meshgrid",
     "nonzero",
+    "ones",
     "place",
     "put",
     "put_along_axis",

dpctl_ext/tensor/_ctors.py

@@ -1518,6 +1518,74 @@ def meshgrid(*arrays, indexing="xy"):
     return output
 
 
+def ones(
+    shape,
+    *,
+    dtype=None,
+    order="C",
+    device=None,
+    usm_type="device",
+    sycl_queue=None,
+):
+    """ones(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 ones.
+
+    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:
+            Created array initialized with ones.
+    """
+    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)
+    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, full_ev = ti._full_usm_ndarray(1, res, sycl_queue)
+    _manager.add_event_pair(hev, full_ev)
+    return res
+
+
 def tril(x, /, *, k=0):
     """
     Returns the lower triangular part of a matrix (or a stack of matrices)

dpnp/dpnp_container.py

@@ -260,7 +260,7 @@ def ones(
         order = "C"
 
     """Creates `dpnp_array` of ones with the given shape, dtype, and order."""
-    array_obj = dpt.ones(
+    array_obj = dpt_ext.ones(
         shape,
         dtype=dtype,
         order=order,

dpnp/dpnp_iface_arraycreation.py

@@ -3696,7 +3696,7 @@ def tri(
     if usm_type is None:
         usm_type = "device"
 
-    m = dpt.ones(
+    m = dpt_ext.ones(
         (N, M),
         dtype=_dtype,
         device=device,

dpnp/tests/test_memory.py

@@ -2,6 +2,9 @@
 import numpy
 import pytest
 
+# TODO: revert to `import dpctl.tensor...`
+# when dpnp fully migrates dpctl/tensor
+import dpctl_ext.tensor as dpt_ext
 import dpnp
 import dpnp.memory as dpm
 
@@ -21,7 +24,7 @@ def test_wrong_input_type(self, x):
             dpm.create_data(x)
 
     def test_wrong_usm_data(self):
-        a = dpt.ones(10)
+        a = dpt_ext.ones(10)
         d = IntUsmData(a.shape, buffer=a)
 
         with pytest.raises(TypeError):

dpnp/tests/test_sycl_queue.py

@@ -2,12 +2,14 @@
 import tempfile
 
 import dpctl
-import dpctl.tensor as dpt
 import numpy
 import pytest
 from dpctl.utils import ExecutionPlacementError
 from numpy.testing import assert_array_equal, assert_raises
 
+# TODO: revert to `import dpctl.tensor...`
+# when dpnp fully migrates dpctl/tensor
+import dpctl_ext.tensor as dpt
 import dpnp
 import dpnp.linalg
 from dpnp.dpnp_array import dpnp_array

dpnp/tests/test_usm_type.py

@@ -2,11 +2,13 @@
 import tempfile
 from math import prod
 
-import dpctl.tensor as dpt
 import dpctl.utils as du
 import numpy
 import pytest
 
+# TODO: revert to `import dpctl.tensor...`
+# when dpnp fully migrates dpctl/tensor
+import dpctl_ext.tensor as dpt
 import dpnp
 from dpnp.dpnp_utils import get_usm_allocations