implementation of dpnp.require (#2036) (#2054)

* add dpnp.require

* improve coverage

* address comments

* fix pre-commit

* Input array should be dpnp.ndarray or usm_ndarray

* update CHANGELOG.md

Author: vtavana

CHANGELOG.md

@@ -39,6 +39,7 @@ In addition, this release completes implementation of `dpnp.fft` module and adds
 * Added runtime dependency on `intel-gpu-ocl-icd-system` package [#2023](https://github.com/IntelPython/dpnp/pull/2023)
 * Added implementation of `dpnp.ravel_multi_index` and `dpnp.unravel_index` functions [#2022](https://github.com/IntelPython/dpnp/pull/2022)
 * Added implementation of `dpnp.resize` and `dpnp.rot90` functions [#2030](https://github.com/IntelPython/dpnp/pull/2030)
+* Added implementation of `dpnp.require` function [#2036](https://github.com/IntelPython/dpnp/pull/2036)
 
 ### Change
 

dpnp/dpnp_iface_manipulation.py

@@ -77,6 +77,7 @@
     "permute_dims",
     "ravel",
     "repeat",
+    "require",
     "reshape",
     "resize",
     "result_type",
@@ -649,12 +650,8 @@ def atleast_1d(*arys):
     """
 
     res = []
+    dpnp.check_supported_arrays_type(*arys)
     for ary in arys:
-        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:
@@ -707,12 +704,8 @@ def atleast_2d(*arys):
     """
 
     res = []
+    dpnp.check_supported_arrays_type(*arys)
     for ary in arys:
-        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, 1)
         elif ary.ndim == 1:
@@ -771,12 +764,8 @@ def atleast_3d(*arys):
     """
 
     res = []
+    dpnp.check_supported_arrays_type(*arys)
     for ary in arys:
-        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, 1, 1)
         elif ary.ndim == 1:
@@ -1959,6 +1948,113 @@ def repeat(a, repeats, axis=None):
     return dpnp_array._create_from_usm_ndarray(usm_res)
 
 
+def require(a, dtype=None, requirements=None, *, like=None):
+    """
+    Return a :class:`dpnp.ndarray` of the provided type that satisfies
+    requirements.
+
+    This function is useful to be sure that an array with the correct flags
+    is returned for passing to compiled code (perhaps through ctypes).
+
+    For full documentation refer to :obj:`numpy.require`.
+
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+       The input array to be converted to a type-and-requirement-satisfying
+       array.
+    dtype : {None, data-type}, optional
+       The required data-type. If ``None`` preserve the current dtype.
+    requirements : {None, str, sequence of str}, optional
+       The requirements list can be any of the following:
+
+       * 'F_CONTIGUOUS' ('F') - ensure a Fortran-contiguous array
+       * 'C_CONTIGUOUS' ('C') - ensure a C-contiguous array
+       * 'WRITABLE' ('W') - ensure a writable array
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        Array with specified requirements and type if given.
+
+    Limitations
+    -----------
+    Parameter `like` is supported only with default value ``None``.
+    Otherwise, the function raises `NotImplementedError` exception.
+
+    See Also
+    --------
+    :obj:`dpnp.asarray` : Convert input to an ndarray.
+    :obj:`dpnp.asanyarray` : Convert to an ndarray, but pass through
+                        ndarray subclasses.
+    :obj:`dpnp.ascontiguousarray` : Convert input to a contiguous array.
+    :obj:`dpnp.asfortranarray` : Convert input to an ndarray with
+                        column-major memory order.
+    :obj:`dpnp.ndarray.flags` : Information about the memory layout
+                        of the array.
+
+    Notes
+    -----
+    The returned array will be guaranteed to have the listed requirements
+    by making a copy if needed.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> x = np.arange(6).reshape(2, 3)
+    >>> x.flags
+      C_CONTIGUOUS : True
+      F_CONTIGUOUS : False
+      WRITEABLE : True
+
+    >>> y = np.require(x, dtype=np.float32, requirements=['W', 'F'])
+    >>> y.flags
+      C_CONTIGUOUS : False
+      F_CONTIGUOUS : True
+      WRITEABLE : True
+
+    """
+
+    dpnp.check_limitations(like=like)
+    dpnp.check_supported_arrays_type(a)
+
+    possible_flags = {
+        "C": "C",
+        "C_CONTIGUOUS": "C",
+        "F": "F",
+        "F_CONTIGUOUS": "F",
+        "W": "W",
+        "WRITEABLE": "W",
+    }
+
+    if not requirements:
+        return dpnp.asanyarray(a, dtype=dtype)
+
+    try:
+        requirements = {possible_flags[x.upper()] for x in requirements}
+    except KeyError as exc:
+        incorrect_flag = (set(requirements) - set(possible_flags.keys())).pop()
+        raise ValueError(
+            f"Incorrect flag {incorrect_flag} in requirements"
+        ) from exc
+
+    order = "A"
+    if requirements.issuperset({"C", "F"}):
+        raise ValueError("Cannot specify both 'C' and 'F' order")
+    if "F" in requirements:
+        order = "F"
+        requirements.remove("F")
+    elif "C" in requirements:
+        order = "C"
+        requirements.remove("C")
+
+    arr = dpnp.array(a, dtype=dtype, order=order, copy=None)
+    if not arr.flags["W"]:
+        return arr.copy(order)
+
+    return arr
+
+
 def reshape(a, /, newshape, order="C", copy=None):
     """
     Gives a new shape to an array without changing its data.

tests/test_arraymanipulation.py

@@ -1,3 +1,4 @@
+import dpctl.tensor as dpt
 import numpy
 import pytest
 from dpctl.tensor._numpy_helper import AxisError
@@ -45,6 +46,13 @@ def test_3D_array(self):
         desired = [a, b]
         assert_array_equal(res, desired)
 
+    def test_dpnp_dpt_array(self):
+        a = dpnp.array([1, 2])
+        b = dpt.asarray([2, 3])
+        res = dpnp.atleast_1d(a, b)
+        desired = [dpnp.array([1, 2]), dpnp.array([2, 3])]
+        assert_array_equal(res, desired)
+
 
 class TestAtleast2d:
     def test_0D_array(self):
@@ -77,6 +85,13 @@ def test_3D_array(self):
         desired = [a, b]
         assert_array_equal(res, desired)
 
+    def test_dpnp_dpt_array(self):
+        a = dpnp.array([1, 2])
+        b = dpt.asarray([2, 3])
+        res = dpnp.atleast_2d(a, b)
+        desired = [dpnp.array([[1, 2]]), dpnp.array([[2, 3]])]
+        assert_array_equal(res, desired)
+
 
 class TestAtleast3d:
     def test_0D_array(self):
@@ -109,6 +124,13 @@ def test_3D_array(self):
         desired = [a, b]
         assert_array_equal(res, desired)
 
+    def test_dpnp_dpt_array(self):
+        a = dpnp.array([1, 2])
+        b = dpt.asarray([2, 3])
+        res = dpnp.atleast_3d(a, b)
+        desired = [dpnp.array([[[1], [2]]]), dpnp.array([[[2], [3]]])]
+        assert_array_equal(res, desired)
+
 
 class TestColumnStack:
     def test_non_iterable(self):

tests/test_manipulation.py

@@ -1,3 +1,5 @@
+import itertools
+
 import dpctl.tensor as dpt
 import numpy
 import pytest
@@ -665,6 +667,73 @@ def test_minimum_signed_integers(self, data, dtype):
         assert_array_equal(result, expected)
 
 
+class TestRequire:
+    flag_names = ["C", "C_CONTIGUOUS", "F", "F_CONTIGUOUS", "W"]
+
+    def generate_all_false(self, dtype):
+        a_np = numpy.zeros((10, 10), dtype=dtype)
+        a_dp = dpnp.zeros((10, 10), dtype=dtype)
+        a_np = a_np[::2, ::2]
+        a_dp = a_dp[::2, ::2]
+        a_np.flags["W"] = False
+        a_dp.flags["W"] = False
+        assert not a_dp.flags["C"]
+        assert not a_dp.flags["F"]
+        assert not a_dp.flags["W"]
+        return a_np, a_dp
+
+    def set_and_check_flag(self, flag, dtype, arr):
+        if dtype is None:
+            dtype = arr[1].dtype
+        result = numpy.require(arr[0], dtype, [flag])
+        expected = dpnp.require(arr[1], dtype, [flag])
+        assert result.flags[flag] == expected.flags[flag]
+        assert result.dtype == expected.dtype
+
+        # a further call to dpnp.require ought to return the same array
+        c = dpnp.require(expected, None, [flag])
+        assert c is expected
+
+    def test_require_each(self):
+        id = ["f4", "i4"]
+        fd = [None, "f4", "c8"]
+        for idtype, fdtype, flag in itertools.product(id, fd, self.flag_names):
+            a = self.generate_all_false(idtype)
+            self.set_and_check_flag(flag, fdtype, a)
+
+    def test_unknown_requirement(self):
+        a = self.generate_all_false("f4")
+        assert_raises(KeyError, numpy.require, a[0], None, "Q")
+        assert_raises(ValueError, dpnp.require, a[1], None, "Q")
+
+    def test_non_array_input(self):
+        a_np = numpy.array([1, 2, 3, 4])
+        a_dp = dpnp.array(a_np)
+        expected = numpy.require(a_np, "i4", ["C", "W"])
+        result = dpnp.require(a_dp, "i4", ["C", "W"])
+        assert expected.flags["C"] == result.flags["C"]
+        assert expected.flags["F"] == result.flags["F"]
+        assert expected.flags["W"] == result.flags["W"]
+        assert expected.dtype == result.dtype
+        assert_array_equal(expected, result)
+
+    def test_C_and_F_simul(self):
+        a = self.generate_all_false("f4")
+        assert_raises(ValueError, numpy.require, a[0], None, ["C", "F"])
+        assert_raises(ValueError, dpnp.require, a[1], None, ["C", "F"])
+
+    def test_copy(self):
+        a_np = numpy.arange(6).reshape(2, 3)
+        a_dp = dpnp.arange(6).reshape(2, 3)
+        a_np.flags["W"] = False
+        a_dp.flags["W"] = False
+        expected = numpy.require(a_np, requirements=["W", "C"])
+        result = dpnp.require(a_dp, requirements=["W", "C"])
+        # copy is done
+        assert result is not a_dp
+        assert_array_equal(expected, result)
+
+
 class TestResize:
     @pytest.mark.parametrize(
         "data, shape",

tests/test_sycl_queue.py

@@ -1285,6 +1285,26 @@ def test_out_multi_dot(device):
         assert_sycl_queue_equal(result.sycl_queue, exec_q)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_require(device):
+    dpnp_data = dpnp.arange(10, device=device).reshape(2, 5)
+    result = dpnp.require(dpnp_data, dtype="f4", requirements=["F"])
+
+    expected_queue = dpnp_data.sycl_queue
+    result_queue = result.sycl_queue
+    assert_sycl_queue_equal(result_queue, expected_queue)
+
+    # No requirements
+    result = dpnp.require(dpnp_data, dtype="f4")
+    expected_queue = dpnp_data.sycl_queue
+    result_queue = result.sycl_queue
+    assert_sycl_queue_equal(result_queue, expected_queue)
+
+
 @pytest.mark.parametrize(
     "device",
     valid_devices,

tests/test_usm_type.py

@@ -1013,6 +1013,19 @@ def test_eigenvalue(func, shape, usm_type):
     assert a.usm_type == dp_val.usm_type
 
 
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+def test_require(usm_type):
+    dpnp_data = dp.arange(10, usm_type=usm_type).reshape(2, 5)
+    result = dp.require(dpnp_data, dtype="f4", requirements=["F"])
+    assert dpnp_data.usm_type == usm_type
+    assert result.usm_type == usm_type
+
+    # No requirements
+    result = dp.require(dpnp_data, dtype="f4")
+    assert dpnp_data.usm_type == usm_type
+    assert result.usm_type == usm_type
+
+
 @pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
 def test_resize(usm_type):
     dpnp_data = dp.arange(10, usm_type=usm_type)

tests/third_party/cupy/manipulation_tests/test_kind.py

@@ -94,7 +94,6 @@ def func(xp):
 
         assert func(numpy) == func(cupy)
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
     @testing.for_all_dtypes()
     def test_require_flag_check(self, dtype):
         possible_flags = [["C_CONTIGUOUS"], ["F_CONTIGUOUS"]]
@@ -105,36 +104,32 @@ def test_require_flag_check(self, dtype):
                 assert arr.flags[parameter]
                 assert arr.dtype == dtype
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
+    @pytest.mark.skip("dpnp.require() does not support requirement ['O']")
     @testing.for_all_dtypes()
     def test_require_owndata(self, dtype):
         x = cupy.zeros((2, 3, 4), dtype=dtype)
         arr = x.view()
         arr = cupy.require(arr, dtype, ["O"])
         assert arr.flags["OWNDATA"]
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
     @testing.for_all_dtypes()
     def test_require_C_and_F_flags(self, dtype):
         x = cupy.zeros((2, 3, 4), dtype=dtype)
         with pytest.raises(ValueError):
             cupy.require(x, dtype, ["C", "F"])
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
     @testing.for_all_dtypes()
     def test_require_incorrect_requirments(self, dtype):
         x = cupy.zeros((2, 3, 4), dtype=dtype)
         with pytest.raises(ValueError):
-            cupy.require(x, dtype, ["W"])
+            cupy.require(x, dtype, ["O"])
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
     @testing.for_all_dtypes()
     def test_require_incorrect_dtype(self, dtype):
         x = cupy.zeros((2, 3, 4), dtype=dtype)
-        with pytest.raises(ValueError):
+        with pytest.raises((ValueError, TypeError)):
             cupy.require(x, "random", "C")
 
-    @pytest.mark.skip("dpnp.require() is not implemented yet")
     @testing.for_all_dtypes()
     def test_require_empty_requirements(self, dtype):
         x = cupy.zeros((2, 3, 4), dtype=dtype)