diff --git a/CHANGELOG.md b/CHANGELOG.md index 201ceb51716..e6f879f575c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -90,6 +90,7 @@ This release is compatible with NumPy 2.5. * Fixed `dpnp.all` and `dpnp.any` aborting when reducing over an empty axis (e.g. an array with a zero-length dimension) [#3021](https://github.com/IntelPython/dpnp/pull/3021) * Released the GIL before the blocking OneMKL DFT calls in the FFT extension [#3040](https://github.com/IntelPython/dpnp/pull/3040) * Fixed `astype` casting an out-of-range floating point value to a signed narrow integer type saturating to the destination min/max instead of wrapping like NumPy, generalizing the earlier unsigned-only fix [#3033](https://github.com/IntelPython/dpnp/pull/3033) +* Fixed `dpnp.ndarray.flat` indexing edge cases, adding support for slices, ellipsis, and integer/boolean array indices [#3045](https://github.com/IntelPython/dpnp/pull/3045) ### Security diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py index 86055a4828f..ea981c6f2c4 100644 --- a/dpnp/dpnp_array.py +++ b/dpnp/dpnp_array.py @@ -1332,9 +1332,37 @@ def flags(self): @property def flat(self): """ - Return a flat iterator, or set a flattened version of self to value. + A 1-D iterator over the array. - """ # noqa: D200 + This is a :obj:`dpnp.flatiter` instance, which acts similarly to, but + is not a subclass of, Python's built-in iterator object. + + For full documentation refer to :obj:`numpy.ndarray.flat`. + + See Also + -------- + :obj:`dpnp.flatiter` : Flat iterator object to iterate over arrays. + :obj:`dpnp.ndarray.flatten` : Return a flattened copy of the array. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(1, 7).reshape(2, 3) + >>> x + array([[1, 2, 3], + [4, 5, 6]]) + >>> x.flat[3] + array(4) + >>> x.T.flat[3] + array(5) + + An assignment example: + + >>> x.flat[[1, 4]] = 1; x + array([[1, 1, 3], + [4, 1, 6]]) + + """ return dpnp.flatiter(self) @@ -1367,7 +1395,7 @@ def flatten(self, /, order="C"): See Also -------- :obj:`dpnp.ravel` : Return a flattened array. - :obj:`dpnp.flat` : A 1-D flat iterator over the array. + :obj:`dpnp.ndarray.flat` : A 1-D flat iterator over the array. Examples -------- diff --git a/dpnp/dpnp_flatiter.py b/dpnp/dpnp_flatiter.py index 7375e03d802..9651057bd90 100644 --- a/dpnp/dpnp_flatiter.py +++ b/dpnp/dpnp_flatiter.py @@ -32,59 +32,139 @@ class flatiter: - """Flat iterator object to iterate over arrays.""" + """ + Flat iterator object to iterate over arrays. - def __init__(self, X): - if type(X) is not dpnp.ndarray: + A flat iterator is returned by :obj:`dpnp.ndarray.flat` for any array. It + allows iterating over the array as if it were a 1-D array, either in a + for-loop or by calling its ``next`` method. + + Iteration is done in row-major, C-style order (the last index varying the + fastest). The iterator can also be indexed using basic slicing or advanced + indexing. + + For full documentation refer to :obj:`numpy.flatiter`. + + See Also + -------- + :obj:`dpnp.ndarray.flat` : Return a flat iterator over an array. + :obj:`dpnp.ndarray.flatten` : Return a flattened copy of an array. + + Examples + -------- + >>> import dpnp as np + >>> x = np.arange(6).reshape(2, 3) + >>> for item in x.flat: + ... print(item) + 0 + 1 + 2 + 3 + 4 + 5 + + >>> x.flat[2:4] + array([2, 3]) + + """ + + def __init__(self, a): + if not isinstance(a, dpnp.ndarray): raise TypeError( - "Argument must be of type dpnp.ndarray, got {}".format(type(X)) + f"An array must be of type dpnp.ndarray, but got {type(a)}" + ) + self._arr = a + self._size = a.size + self._i = 0 + + @staticmethod + def _reject_newaxis(key): + # newaxis (None) is valid for array indexing but not for flat indexing + if key is None or ( + isinstance(key, tuple) and any(k is None for k in key) + ): + raise IndexError( + "only integers, slices (`:`), ellipsis (`...`) and integer " + "or boolean arrays are valid indices" ) - self.arr_ = X - self.size_ = X.size - self.i_ = 0 - - def _multiindex(self, i): - nd = self.arr_.ndim - if nd == 0: - if i == 0: - return () - raise KeyError - elif nd == 1: - return (i,) - sh = self.arr_.shape - i_ = i - multi_index = [0] * nd - for k in reversed(range(1, nd)): - si = sh[k] - q = i_ // si - multi_index[k] = i_ - q * si - i_ = q - multi_index[0] = i_ - return tuple(multi_index) + + def _check_bounds(self, key): + # fancy int indices wrap instead of raising, so check them vs NumPy + if key is Ellipsis or isinstance(key, (slice, bool, tuple)): + return + + if isinstance(key, int) or ( + callable(getattr(key, "__index__", None)) + and not hasattr(key, "ndim") + ): + return # scalar int: regular indexing checks it + + try: + idx = dpnp.asarray(key, sycl_queue=self._arr.sycl_queue) + except Exception: + return # let regular indexing raise + + if idx.dtype.kind not in "iu" or idx.size == 0: + return + + size = self._size + hi, lo = int(dpnp.max(idx)), int(dpnp.min(idx)) + if hi >= size: + raise IndexError(f"index {hi} is out of bounds for size {size}") + if lo < -size: + raise IndexError(f"index {lo} is out of bounds for size {size}") + + def _flatten(self): + # C-order flat view (copy if non-contiguous) + return dpnp.reshape(self._arr, -1) def __getitem__(self, key): - idx = getattr(key, "__index__", None) - if not callable(idx): - raise TypeError(key) - i = idx() - mi = self._multiindex(i) - return self.arr_.__getitem__(mi) + self._reject_newaxis(key) + self._check_bounds(key) + + # flat always yields a copy, never a view + return self._flatten()[key].copy() def __setitem__(self, key, val): - idx = getattr(key, "__index__", None) - if not callable(idx): - raise TypeError(key) - i = idx() - mi = self._multiindex(i) - return self.arr_.__setitem__(mi, val) + self._reject_newaxis(key) + self._check_bounds(key) + + if isinstance(key, tuple) and len(key) == 0: + # NumPy rejects arr.flat[()] = val + raise IndexError( + "Assigning to a flat iterator with a 0-D index is not " + "supported" + ) + + a = self._arr + exec_q = a.sycl_queue + usm_type = a.usm_type + + # resolve key to flat positions, reusing regular indexing to validate + flat_index = dpnp.arange(a.size, sycl_queue=exec_q, usm_type=usm_type) + idx = flat_index[key] + + if not dpnp.isscalar(val): + val = dpnp.asarray( + val, sycl_queue=exec_q, usm_type=usm_type + ).ravel() + n = idx.size + if 0 < val.size != n: + # cycles the values over the selection + val = val[ + dpnp.arange(n, sycl_queue=exec_q, usm_type=usm_type) + % val.size + ] + + dpnp.put(a, idx, val) def __iter__(self): return self def __next__(self): - if self.i_ < self.size_: - val = self.__getitem__(self.i_) - self.i_ = self.i_ + 1 + if self._i < self._size: + val = self.__getitem__(self._i) + self._i = self._i + 1 return val else: raise StopIteration diff --git a/dpnp/tests/test_flat.py b/dpnp/tests/test_flat.py index c40e95d3ee8..70bbd1ef335 100644 --- a/dpnp/tests/test_flat.py +++ b/dpnp/tests/test_flat.py @@ -1,9 +1,11 @@ import numpy as np import pytest -from numpy.testing import assert_array_equal, assert_raises +from numpy.testing import assert_array_equal import dpnp +from .third_party.cupy import testing + class TestFlatiter: @pytest.mark.parametrize( @@ -16,37 +18,144 @@ class TestFlatiter: ids=["1D array", "2D array", "2D.T array"], ) def test_flat_getitem(self, a, index): - a_dp = dpnp.array(a) - result = a_dp.flat[index] + ia = dpnp.array(a) + result = ia.flat[index] expected = a.flat[index] assert_array_equal(expected, result) def test_flat_iteration(self): a = np.array([[1, 2], [3, 4]]) - a_dp = dpnp.array(a) - for dp_val, np_val in zip(a_dp.flat, a.flat): - assert dp_val == np_val + ia = dpnp.array(a) + for ival, val in zip(ia.flat, a.flat): + assert ival == val def test_init_error(self): - assert_raises(TypeError, dpnp.flatiter, [1, 2, 3]) + with pytest.raises(TypeError, match="must be of type dpnp.ndarray"): + dpnp.flatiter([1, 2, 3]) + + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_key_error(self, xp): + a = xp.array(42) + with pytest.raises(IndexError): + _ = a.flat[1] - def test_flat_key_error(self): - a_dp = dpnp.array(42) - with pytest.raises(KeyError): - _ = a_dp.flat[1] + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_invalid_key(self, xp): + flat = xp.array([1, 2, 3]).flat - def test_flat_invalid_key(self): - a_dp = dpnp.array([1, 2, 3]) - flat = dpnp.flatiter(a_dp) # check __getitem__ - with pytest.raises(TypeError): + with pytest.raises(IndexError): _ = flat["invalid"] + # check __setitem__ - with pytest.raises(TypeError): + with pytest.raises(IndexError): flat["invalid"] = 42 - def test_flat_out_of_bounds(self): - a_dp = dpnp.array([1, 2, 3]) - flat = dpnp.flatiter(a_dp) + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_out_of_bounds(self, xp): + flat = xp.array([1, 2, 3]).flat with pytest.raises(IndexError): _ = flat[10] + + @pytest.mark.parametrize( + "key", + [ + slice(1, 4), + slice(None), + slice(None, None, 2), + [0, 2, 4], + [-1, -2], + Ellipsis, + ], + ids=["slice", "full_slice", "step_slice", "list", "neg_list", "..."], + ) + def test_flat_getitem_index_types(self, key): + a = np.arange(1, 7).reshape(2, 3) + ia = dpnp.array(a) + assert_array_equal(ia.flat[key], a.flat[key]) + + @pytest.mark.parametrize( + "key", + [slice(1, 4), slice(None), [0, 2, 4], [-1, -2], Ellipsis], + ids=["slice", "full_slice", "list", "neg_list", "..."], + ) + def test_flat_setitem_index_types(self, key): + a = np.arange(1, 7).reshape(2, 3) + ia = dpnp.array(a) + a.flat[key] = 0 + ia.flat[key] = 0 + assert_array_equal(ia, a) + + def test_flat_index_array(self): + a = np.arange(1, 7).reshape(2, 3) + ia = dpnp.array(a) + + # int array index + assert_array_equal(ia.flat[dpnp.array([0, 3, 5])], a.flat[[0, 3, 5]]) + + @testing.with_requires("numpy>=2.4") + def test_flat_bool_mask(self): + a = np.arange(1, 7).reshape(2, 3) + ia = dpnp.array(a) + mask = np.array([True, False] * 3) + + # getitem via bool array + assert_array_equal(ia.flat[dpnp.array(mask)], a.flat[mask]) + + # setitem via bool array + a.flat[mask] = -1 + ia.flat[dpnp.array(mask)] = -1 + assert_array_equal(ia, a) + + def test_flat_non_contiguous(self): + # C-order traversal + write-back for non-contiguous arrays + a = np.arange(1, 7).reshape(2, 3).T + ia = dpnp.array(np.arange(1, 7).reshape(2, 3)).T + assert_array_equal(ia.flat[1:5], a.flat[1:5]) + a.flat[1:5] = 0 + ia.flat[1:5] = 0 + assert_array_equal(ia, a) + + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_getitem_returns_copy(self, xp): + # flat yields copies, not views + a = xp.arange(10) + s = a.flat[1:4] + s[0] = 999 + assert a[1] != 999 + + def test_flat_scalar_getitem_returns_copy(self): + # dpnp returns a 0-d array copy (NumPy returns an immutable scalar) + ia = dpnp.arange(10) + x = ia.flat[3] + x[...] = 777 + assert ia[3] != 777 + + @testing.with_requires("numpy>=2.4") + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_newaxis(self, xp): + a = xp.array([1, 2, 3]) + with pytest.raises(IndexError, match="are valid indices"): + _ = a.flat[None] + with pytest.raises(IndexError, match="are valid indices"): + a.flat[None] = 0 + + @testing.with_requires("numpy>=2.4") + @pytest.mark.parametrize("xp", [dpnp, np]) + def test_flat_empty_tuple(self, xp): + a = xp.arange(1, 7).reshape(2, 3) + # getitem with () returns the whole flattened array + assert_array_equal(a.flat[()], xp.arange(1, 7)) + # setitem with a 0-d index is unsupported + with pytest.raises(IndexError, match="0-D index is not supported"): + a.flat[()] = 0 + + @testing.with_requires("numpy>=2.4") + @pytest.mark.parametrize("xp", [dpnp, np]) + @pytest.mark.parametrize("key", [[100], [-100]], ids=["oob", "neg_oob"]) + def test_flat_array_out_of_bounds(self, xp, key): + a = xp.array([1, 2, 3]) + with pytest.raises(IndexError, match="out of bounds for size"): + _ = a.flat[key] + with pytest.raises(IndexError, match="out of bounds for size"): + a.flat[key] = 0 diff --git a/dpnp/tests/third_party/cupy/indexing_tests/test_iterate.py b/dpnp/tests/third_party/cupy/indexing_tests/test_iterate.py index f68af146dd6..7cf8d995b4b 100644 --- a/dpnp/tests/third_party/cupy/indexing_tests/test_iterate.py +++ b/dpnp/tests/third_party/cupy/indexing_tests/test_iterate.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import unittest import warnings @@ -58,17 +60,17 @@ def test_copy_next(self, xp): @testing.parameterize( - # {"shape": (2, 3, 4), "index": Ellipsis}, + {"shape": (2, 3, 4), "index": Ellipsis}, {"shape": (2, 3, 4), "index": 0}, {"shape": (2, 3, 4), "index": 10}, - # {"shape": (2, 3, 4), "index": slice(None)}, - # {"shape": (2, 3, 4), "index": slice(None, 10)}, - # {"shape": (2, 3, 4), "index": slice(None, None, 2)}, - # {"shape": (2, 3, 4), "index": slice(None, None, -1)}, - # {"shape": (2, 3, 4), "index": slice(10, None, -1)}, - # {"shape": (2, 3, 4), "index": slice(10, None, -2)}, - # {"shape": (), "index": slice(None)}, - # {"shape": (10,), "index": slice(None)}, + {"shape": (2, 3, 4), "index": slice(None)}, + {"shape": (2, 3, 4), "index": slice(None, 10)}, + {"shape": (2, 3, 4), "index": slice(None, None, 2)}, + {"shape": (2, 3, 4), "index": slice(None, None, -1)}, + {"shape": (2, 3, 4), "index": slice(10, None, -1)}, + {"shape": (2, 3, 4), "index": slice(10, None, -2)}, + {"shape": (), "index": slice(None)}, + {"shape": (10,), "index": slice(None)}, ) class TestFlatiterSubscript(unittest.TestCase): @@ -125,12 +127,13 @@ def test_setitem_ndarray_different_types(self, xp, a_dtype, v_dtype, order): @testing.parameterize( {"shape": (2, 3, 4), "index": None}, - {"shape": (2, 3, 4), "index": (0,)}, - {"shape": (2, 3, 4), "index": True}, - {"shape": (2, 3, 4), "index": cupy.array([0])}, - {"shape": (2, 3, 4), "index": [0]}, + # the indices below are valid for flat iterators since NumPy 2.4 + # (numpy-gh-28590) and no longer raise an IndexError: + # {"shape": (2, 3, 4), "index": (0,)}, + # {"shape": (2, 3, 4), "index": True}, + # {"shape": (2, 3, 4), "index": cupy.array([0])}, + # {"shape": (2, 3, 4), "index": [0]}, ) -@pytest.mark.skip("no exception raised") class TestFlatiterSubscriptIndexError(unittest.TestCase): @testing.for_all_dtypes()