Merge pull request #1585 from IntelPython/update_atleast_func

Updated atleast_2d and atleast_3d functions.

Author: npolina4

dpnp/dpnp_algo/dpnp_algo_manipulation.pxi

@@ -36,8 +36,6 @@ and the rest of the library
 # NO IMPORTs here. All imports must be placed into main "dpnp_algo.pyx" file
 
 __all__ += [
-    "dpnp_atleast_2d",
-    "dpnp_atleast_3d",
     "dpnp_repeat",
 ]
 
@@ -48,60 +46,6 @@ ctypedef c_dpctl.DPCTLSyclEventRef(*fptr_dpnp_repeat_t)(c_dpctl.DPCTLSyclQueueRe
                                                         const c_dpctl.DPCTLEventVectorRef)
 
 
-cpdef utils.dpnp_descriptor dpnp_atleast_2d(utils.dpnp_descriptor arr):
-    # it looks like it should be dpnp.copy + dpnp.reshape
-    cdef utils.dpnp_descriptor result
-    cdef size_t arr_ndim = arr.ndim
-    cdef long arr_size = arr.size
-    if arr_ndim == 1:
-        arr_obj = arr.get_array()
-        result = utils_py.create_output_descriptor_py((1, arr_size),
-                                                      arr.dtype,
-                                                      None,
-                                                      device=arr_obj.sycl_device,
-                                                      usm_type=arr_obj.usm_type,
-                                                      sycl_queue=arr_obj.sycl_queue)
-        for i in range(arr_size):
-            result.get_pyobj()[0, i] = arr.get_pyobj()[i]
-        return result
-    else:
-        return arr
-
-
-cpdef utils.dpnp_descriptor dpnp_atleast_3d(utils.dpnp_descriptor arr):
-    # it looks like it should be dpnp.copy + dpnp.reshape
-    cdef utils.dpnp_descriptor result
-    cdef size_t arr_ndim = arr.ndim
-    cdef shape_type_c arr_shape = arr.shape
-    cdef long arr_size = arr.size
-
-    arr_obj = arr.get_array()
-
-    if arr_ndim == 1:
-        result = utils_py.create_output_descriptor_py((1, 1, arr_size),
-                                                      arr.dtype,
-                                                      None,
-                                                      device=arr_obj.sycl_device,
-                                                      usm_type=arr_obj.usm_type,
-                                                      sycl_queue=arr_obj.sycl_queue)
-        for i in range(arr_size):
-            result.get_pyobj()[0, 0, i] = arr.get_pyobj()[i]
-        return result
-    elif arr_ndim == 2:
-        result = utils_py.create_output_descriptor_py((1, arr_shape[0], arr_shape[1]),
-                                                      arr.dtype,
-                                                      None,
-                                                      device=arr_obj.sycl_device,
-                                                      usm_type=arr_obj.usm_type,
-                                                      sycl_queue=arr_obj.sycl_queue)
-        for i in range(arr_shape[0]):
-            for j in range(arr_shape[1]):
-                result.get_pyobj()[0, i, j] = arr.get_pyobj()[i, j]
-        return result
-    else:
-        return arr
-
-
 cpdef utils.dpnp_descriptor dpnp_repeat(utils.dpnp_descriptor array1, repeats, axes=None):
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(array1.dtype)
 

dpnp/dpnp_iface_manipulation.py

@@ -124,15 +124,13 @@ def atleast_1d(*arys):
     """
     Convert inputs to arrays with at least one dimension.
 
-    Scalar inputs are converted to 1-dimensional arrays, whilst
-    higher-dimensional inputs are preserved.
-
     For full documentation refer to :obj:`numpy.atleast_1d`.
 
     Parameters
     ----------
-    arys : {dpnp_array, usm_ndarray}
-        One or more input arrays.
+    arys : {dpnp.ndarray, usm_ndarray}
+        One or more array-like sequences. Arrays that already have one or more
+        dimensions are preserved.
 
     Returns
     -------
@@ -142,14 +140,20 @@ def atleast_1d(*arys):
 
     See Also
     --------
-    atleast_2d, atleast_3d
+    :obj:`dpnp.atleast_2d` : View inputs as arrays with at least two dimensions.
+    :obj:`dpnp.atleast_3d` : View inputs as arrays with at least three dimensions.
 
     Examples
     --------
     >>> import dpnp as np
-    >>> np.atleast_1d(1.0)
+    >>> x = np.array(1.0)
+    >>> np.atleast_1d(x)
     array([1.])
 
+    >>> y = np.array([3, 4])
+    >>> np.atleast_1d(x, y)
+    [array([1.]), array([3, 4])]
+
     >>> x = np.arange(9.0).reshape(3,3)
     >>> np.atleast_1d(x)
     array([[0., 1., 2.],
@@ -158,18 +162,21 @@ def atleast_1d(*arys):
     >>> np.atleast_1d(x) is x
     True
 
-    >>> np.atleast_1d(1, [3, 4])
-    [array([1]), array([3, 4])]
-
     """
 
     res = []
     for ary in arys:
-        ary = dpnp.asanyarray(ary)
+        if not dpnp.is_supported_array_type(ary):
+            raise TypeError(
+                "Each input array must be any of supported type, "
+                f"but got {type(ary)}"
+            )
         if ary.ndim == 0:
             result = ary.reshape(1)
         else:
             result = ary
+        if isinstance(result, dpt.usm_ndarray):
+            result = dpnp_array._create_from_usm_ndarray(result)
         res.append(result)
     if len(res) == 1:
         return res[0]
@@ -183,36 +190,57 @@ def atleast_2d(*arys):
 
     For full documentation refer to :obj:`numpy.atleast_2d`.
 
-    Limitations
-    -----------
-    Input arrays is supported as :obj:`dpnp.ndarray`.
+    Parameters
+    ----------
+    arys : {dpnp.ndarray, usm_ndarray}
+        One or more array-like sequences. Arrays that already have two or more
+        dimensions are preserved.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        An array, or list of arrays, each with ``a.ndim >= 2``.
+        Copies are avoided where possible, and views with two or more
+        dimensions are returned.
+
+    See Also
+    --------
+    :obj:`dpnp.atleast_1d` : Convert inputs to arrays with at least one dimension.
+    :obj:`dpnp.atleast_3d` : View inputs as arrays with at least three dimensions.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> x = np.array(3.0)
+    >>> np.atleast_2d(x)
+    array([[3.]])
+
+    >>> x = np.arange(3.0)
+    >>> np.atleast_2d(x)
+    array([[0., 1., 2.]])
+
     """
 
-    all_is_array = True
-    arys_desc = []
+    res = []
     for ary in arys:
-        if not dpnp.isscalar(ary):
-            ary_desc = dpnp.get_dpnp_descriptor(
-                ary, copy_when_nondefault_queue=False
+        if not dpnp.is_supported_array_type(ary):
+            raise TypeError(
+                "Each input array must be any of supported type, "
+                f"but got {type(ary)}"
             )
-            if ary_desc:
-                arys_desc.append(ary_desc)
-                continue
-        all_is_array = False
-        break
-
-    if not use_origin_backend(arys[0]) and all_is_array:
-        result = []
-        for ary_desc in arys_desc:
-            res = dpnp_atleast_2d(ary_desc).get_pyobj()
-            result.append(res)
-
-        if len(result) == 1:
-            return result[0]
+        if ary.ndim == 0:
+            result = ary.reshape(1, 1)
+        elif ary.ndim == 1:
+            result = ary[dpnp.newaxis, :]
         else:
-            return result
-
-    return call_origin(numpy.atleast_2d, *arys)
+            result = ary
+        if isinstance(result, dpt.usm_ndarray):
+            result = dpnp_array._create_from_usm_ndarray(result)
+        res.append(result)
+    if len(res) == 1:
+        return res[0]
+    else:
+        return res
 
 
 def atleast_3d(*arys):
@@ -221,36 +249,63 @@ def atleast_3d(*arys):
 
     For full documentation refer to :obj:`numpy.atleast_3d`.
 
-    Limitations
-    -----------
-    Input arrays is supported as :obj:`dpnp.ndarray`.
+    Parameters
+    ----------
+    arys : {dpnp.ndarray, usm_ndarray}
+        One or more array-like sequences. Arrays that already have three or more
+        dimensions are preserved.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        An array, or list of arrays, each with ``a.ndim >= 3``. Copies are
+        avoided where possible, and views with three or more dimensions are
+        returned.
+
+    See Also
+    --------
+    :obj:`dpnp.atleast_1d` : Convert inputs to arrays with at least one dimension.
+    :obj:`dpnp.atleast_2d` : View inputs as arrays with at least three dimensions.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> x = np.array(3.0)
+    >>> np.atleast_3d(x)
+    array([[[3.]]])
+
+    >>> x = np.arange(3.0)
+    >>> np.atleast_3d(x).shape
+    (1, 3, 1)
+
+    >>> x = np.arange(12.0).reshape(4, 3)
+    >>> np.atleast_3d(x).shape
+    (4, 3, 1)
+
     """
 
-    all_is_array = True
-    arys_desc = []
+    res = []
     for ary in arys:
-        if not dpnp.isscalar(ary):
-            ary_desc = dpnp.get_dpnp_descriptor(
-                ary, copy_when_nondefault_queue=False
+        if not dpnp.is_supported_array_type(ary):
+            raise TypeError(
+                "Each input array must be any of supported type, "
+                f"but got {type(ary)}"
             )
-            if ary_desc:
-                arys_desc.append(ary_desc)
-                continue
-        all_is_array = False
-        break
-
-    if not use_origin_backend(arys[0]) and all_is_array:
-        result = []
-        for ary_desc in arys_desc:
-            res = dpnp_atleast_3d(ary_desc).get_pyobj()
-            result.append(res)
-
-        if len(result) == 1:
-            return result[0]
+        if ary.ndim == 0:
+            result = ary.reshape(1, 1, 1)
+        elif ary.ndim == 1:
+            result = ary[dpnp.newaxis, :, dpnp.newaxis]
+        elif ary.ndim == 2:
+            result = ary[:, :, dpnp.newaxis]
         else:
-            return result
-
-    return call_origin(numpy.atleast_3d, *arys)
+            result = ary
+        if isinstance(result, dpt.usm_ndarray):
+            result = dpnp_array._create_from_usm_ndarray(result)
+        res.append(result)
+    if len(res) == 1:
+        return res[0]
+    else:
+        return res
 
 
 def broadcast_to(array, /, shape, subok=False):

tests/skipped_tests.tbl

@@ -427,6 +427,7 @@ tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_arra
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_diff_length
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_is_equal
 tests/third_party/cupy/logic_tests/test_comparison.py::TestArrayEqual::test_array_equal_not_equal
+
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_0_{shapes=[(), ()]}::test_broadcast_arrays
 tests/third_party/cupy/manipulation_tests/test_dims.py::TestBroadcast_param_10_{shapes=[(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]}::test_broadcast

tests/test_arraymanipulation.py

@@ -655,12 +655,6 @@ def test_3D_array(self):
         desired = [a, b]
         assert_array_equal(res, desired)
 
-    def test_r1array(self):
-        assert dpnp.atleast_1d(3).shape == (1,)
-        assert dpnp.atleast_1d(3j).shape == (1,)
-        assert dpnp.atleast_1d(3.0).shape == (1,)
-        assert dpnp.atleast_1d([[2, 3], [4, 5]]).shape == (2, 2)
-
 
 class TestRollaxis:
     data = [
@@ -725,3 +719,67 @@ def test_results(self):
             jp = j + 1 if j < 4 else j
             res = dpnp.rollaxis(dp_a, axis=-ip, start=-jp)
             exp = numpy.rollaxis(np_a, axis=-ip, start=-jp)
+
+
+class TestAtleast2d:
+    def test_0D_array(self):
+        a = dpnp.array(1)
+        b = dpnp.array(2)
+        res = [dpnp.atleast_2d(a), dpnp.atleast_2d(b)]
+        desired = [dpnp.array([[1]]), dpnp.array([[2]])]
+        assert_array_equal(res, desired)
+
+    def test_1D_array(self):
+        a = dpnp.array([1, 2])
+        b = dpnp.array([2, 3])
+        res = [dpnp.atleast_2d(a), dpnp.atleast_2d(b)]
+        desired = [dpnp.array([[1, 2]]), dpnp.array([[2, 3]])]
+        assert_array_equal(res, desired)
+
+    def test_2D_array(self):
+        a = dpnp.array([[1, 2], [1, 2]])
+        b = dpnp.array([[2, 3], [2, 3]])
+        res = [dpnp.atleast_2d(a), dpnp.atleast_2d(b)]
+        desired = [a, b]
+        assert_array_equal(res, desired)
+
+    def test_3D_array(self):
+        a = dpnp.array([[1, 2], [1, 2]])
+        b = dpnp.array([[2, 3], [2, 3]])
+        a = dpnp.array([a, a])
+        b = dpnp.array([b, b])
+        res = [dpnp.atleast_2d(a), dpnp.atleast_2d(b)]
+        desired = [a, b]
+        assert_array_equal(res, desired)
+
+
+class TestAtleast3d:
+    def test_0D_array(self):
+        a = dpnp.array(1)
+        b = dpnp.array(2)
+        res = [dpnp.atleast_3d(a), dpnp.atleast_3d(b)]
+        desired = [dpnp.array([[[1]]]), dpnp.array([[[2]]])]
+        assert_array_equal(res, desired)
+
+    def test_1D_array(self):
+        a = dpnp.array([1, 2])
+        b = dpnp.array([2, 3])
+        res = [dpnp.atleast_3d(a), dpnp.atleast_3d(b)]
+        desired = [dpnp.array([[[1], [2]]]), dpnp.array([[[2], [3]]])]
+        assert_array_equal(res, desired)
+
+    def test_2D_array(self):
+        a = dpnp.array([[1, 2], [1, 2]])
+        b = dpnp.array([[2, 3], [2, 3]])
+        res = [dpnp.atleast_3d(a), dpnp.atleast_3d(b)]
+        desired = [a[:, :, dpnp.newaxis], b[:, :, dpnp.newaxis]]
+        assert_array_equal(res, desired)
+
+    def test_3D_array(self):
+        a = dpnp.array([[1, 2], [1, 2]])
+        b = dpnp.array([[2, 3], [2, 3]])
+        a = dpnp.array([a, a])
+        b = dpnp.array([b, b])
+        res = [dpnp.atleast_3d(a), dpnp.atleast_3d(b)]
+        desired = [a, b]
+        assert_array_equal(res, desired)