diff --git a/dpnp/dpnp_iface_arraycreation.py b/dpnp/dpnp_iface_arraycreation.py index 28ad6b0b1fc..d33ed6ef6d4 100644 --- a/dpnp/dpnp_iface_arraycreation.py +++ b/dpnp/dpnp_iface_arraycreation.py @@ -866,9 +866,14 @@ def diag(v, /, k=0, *, device=None, usm_type=None, sycl_queue=None): v : array_like Input data, in any form that can be converted to an array. This includes scalars, lists, lists of tuples, tuples, tuples of tuples, - tuples of lists, and ndarrays. If `v` is a 2-D array, return a copy of - its k-th diagonal. If `v` is a 1-D array, return a 2-D array with `v` - on the k-th diagonal. + tuples of lists, and ndarrays. + If `v` is a 1-D array, return a 2-D array with `v` + on the `k`-th diagonal. + If `v` is a 2-D array and is an instance of + {dpnp.ndarray, usm_ndarray}, then: + - If `device`, `usm_type`, and `sycl_queue` are set to their + default values, returns a read/write view of its k-th diagonal. + - Otherwise, returns a copy of its k-th diagonal. k : int, optional Diagonal in question. The default is 0. Use k > 0 for diagonals above the main diagonal, and k < 0 for diagonals below the main diagonal. @@ -894,43 +899,43 @@ def diag(v, /, k=0, *, device=None, usm_type=None, sycl_queue=None): -------- :obj:`diagonal` : Return specified diagonals. :obj:`diagflat` : Create a 2-D array with the flattened input as a diagonal. - :obj:`trace` : Return sum along diagonals. - :obj:`triu` : Return upper triangle of an array. - :obj:`tril` : Return lower triangle of an array. + :obj:`trace` : Return the sum along diagonals of the array. + :obj:`triu` : Upper triangle of an array. + :obj:`tril` : Lower triangle of an array. Examples -------- >>> import dpnp as np - >>> x0 = np.arange(9).reshape((3, 3)) - >>> x0 + >>> x = np.arange(9).reshape((3, 3)) + >>> x array([[0, 1, 2], [3, 4, 5], [6, 7, 8]]) - >>> np.diag(x0) + >>> np.diag(x) array([0, 4, 8]) - >>> np.diag(x0, k=1) + >>> np.diag(x, k=1) array([1, 5]) - >>> np.diag(x0, k=-1) + >>> np.diag(x, k=-1) array([3, 7]) - >>> np.diag(np.diag(x0)) + >>> np.diag(np.diag(x)) array([[0, 0, 0], [0, 4, 0], [0, 0, 8]]) Creating an array on a different device or with a specified usm_type - >>> x = np.diag(x0) # default case - >>> x, x.device, x.usm_type + >>> res = np.diag(x) # default case + >>> res, res.device, res.usm_type (array([0, 4, 8]), Device(level_zero:gpu:0), 'device') - >>> y = np.diag(x0, device="cpu") - >>> y, y.device, y.usm_type + >>> res_cpu = np.diag(x, device="cpu") + >>> res_cpu, res_cpu.device, res_cpu.usm_type (array([0, 4, 8]), Device(opencl:cpu:0), 'device') - >>> z = np.diag(x0, usm_type="host") - >>> z, z.device, z.usm_type + >>> res_host = np.diag(x, usm_type="host") + >>> res_host, res_host.device, res_host.usm_type (array([0, 4, 8]), Device(level_zero:gpu:0), 'host') """ @@ -938,35 +943,18 @@ def diag(v, /, k=0, *, device=None, usm_type=None, sycl_queue=None): if not isinstance(k, int): raise TypeError(f"An integer is required, but got {type(k)}") - v = dpnp.asarray(v, device=device, usm_type=usm_type, sycl_queue=sycl_queue) + v = dpnp.asanyarray( + v, device=device, usm_type=usm_type, sycl_queue=sycl_queue + ) - init0 = max(0, -k) - init1 = max(0, k) if v.ndim == 1: size = v.shape[0] + abs(k) - m = dpnp.zeros( - (size, size), - dtype=v.dtype, - usm_type=v.usm_type, - sycl_queue=v.sycl_queue, - ) - for i in range(v.shape[0]): - m[(init0 + i), init1 + i] = v[i] - return m + ret = dpnp.zeros_like(v, shape=(size, size)) + ret.diagonal(k)[:] = v + return ret if v.ndim == 2: - size = max( - 0, min(v.shape[0], v.shape[0] + k, v.shape[1], v.shape[1] - k) - ) - m = dpnp.zeros( - (size,), - dtype=v.dtype, - usm_type=v.usm_type, - sycl_queue=v.sycl_queue, - ) - for i in range(size): - m[i] = v[(init0 + i), init1 + i] - return m + return v.diagonal(k) raise ValueError("Input must be a 1-D or 2-D array.") @@ -1008,9 +996,9 @@ def diagflat(v, /, k=0, *, device=None, usm_type=None, sycl_queue=None): See Also -------- - :obj:`diag` : Return the extracted diagonal or constructed diagonal array. - :obj:`diagonal` : Return specified diagonals. - :obj:`trace` : Return sum along diagonals. + :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array. + :obj:`dpnp.diagonal` : Return specified diagonals. + :obj:`dpnp.trace` : Return sum along diagonals. Examples -------- @@ -1324,6 +1312,11 @@ def eye( Parameter `like` is supported only with default value ``None``. Otherwise, the function raises `NotImplementedError` exception. + See Also + -------- + :obj:`dpnp.identity` : Return the identity array. + :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array. + Examples -------- >>> import dpnp as np @@ -2264,7 +2257,7 @@ def identity( :obj:`dpnp.eye` : Return a 2-D array with ones on the diagonal and zeros elsewhere. :obj:`dpnp.ones` : Return a new array setting values to one. - :obj:`dpnp.diag` : Return diagonal 2-D array from an input 1-D array. + :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array. Examples -------- diff --git a/tests/third_party/cupy/creation_tests/test_matrix.py b/tests/third_party/cupy/creation_tests/test_matrix.py index d803c6e2850..eda1a40b784 100644 --- a/tests/third_party/cupy/creation_tests/test_matrix.py +++ b/tests/third_party/cupy/creation_tests/test_matrix.py @@ -27,35 +27,35 @@ def test_diag3(self, xp): def test_diag_extraction_from_nested_list(self, xp): a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] r = xp.diag(a, 1) - self.assertIsInstance(r, xp.ndarray) + assert isinstance(r, xp.ndarray) return r @testing.numpy_cupy_array_equal() def test_diag_extraction_from_nested_tuple(self, xp): a = ((1, 2, 3), (4, 5, 6), (7, 8, 9)) r = xp.diag(a, -1) - self.assertIsInstance(r, xp.ndarray) + assert isinstance(r, xp.ndarray) return r @testing.numpy_cupy_array_equal() def test_diag_construction(self, xp): a = testing.shaped_arange((3,), xp) r = xp.diag(a) - self.assertIsInstance(r, xp.ndarray) + assert isinstance(r, xp.ndarray) return r @testing.numpy_cupy_array_equal() def test_diag_construction_from_list(self, xp): a = [1, 2, 3] r = xp.diag(a) - self.assertIsInstance(r, xp.ndarray) + assert isinstance(r, xp.ndarray) return r @testing.numpy_cupy_array_equal() def test_diag_construction_from_tuple(self, xp): a = (1, 2, 3) r = xp.diag(a) - self.assertIsInstance(r, xp.ndarray) + assert isinstance(r, xp.ndarray) return r def test_diag_scaler(self):