Fix dpnp.ndarray.flat indexing edge cases

CHANGELOG.md

@@ -99,6 +99,7 @@ This release is compatible with NumPy 2.5.
 * Fixed `dpnp.cumsum`, `dpnp.cumprod`, and their `nan`/`cumulative_*` variants (including `dpnp.tensor.cumulative_sum`/`cumulative_prod`) silently returning incorrect results when accumulating along an axis of an array with more than one row [#3063](https://github.com/IntelPython/dpnp/pull/3063)
 * Fixed `dpnp.einsum` returning a result whose memory layout differs from NumPy for the default `order="K"`, and ignoring `out` and `order` for a contraction over a size-0 dimension [#3058](https://github.com/IntelPython/dpnp/pull/3058)
 * Fixed operations on a boolean array whose bytes are not `0x00`/`0x01` [#3055](https://github.com/IntelPython/dpnp/pull/3055)
+* Fixed `dpnp.ndarray.flat` indexing and assignment edge cases, adding support for slices, ellipsis, and integer/boolean array indices [#3045](https://github.com/IntelPython/dpnp/pull/3045)
 
 ### Security
 

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
         --------

dpnp/dpnp_flatiter.py

@@ -28,63 +28,232 @@
 
 """Implementation of flatiter."""
 
+import numpy
+
 import dpnp
+import dpnp.tensor as dpt
+
+from .dpnp_array import dpnp_array
 
 
 class flatiter:
-    """Flat iterator object to iterate over arrays."""
+    """
+    Flat iterator object to iterate over arrays.
+
+    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, X):
-        if type(X) is not dpnp.ndarray:
+    def __init__(self, a):
+        if not isinstance(a, dpnp_array):
             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_ = 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)
+        self._arr = a
+        self._size = a.size
+        self._i = 0
+
+    def _validate_key(self, key):
+        """
+        Validate `key` as a flat iterator index.
+
+        Return the array of flat positions for an integer-array key, or
+        ``None`` when the caller has to resolve the positions itself.
+
+        """
+        # Ellipsis/slice/tuple need no validation here
+        if key is Ellipsis or isinstance(key, (slice, tuple)):
+            return None
+
+        # a genuine scalar int (not bool, not an array): bounds-checked later
+        if (
+            not isinstance(key, bool)
+            and callable(getattr(key, "__index__", None))
+            and not hasattr(key, "ndim")
+        ):
+            return None
+
+        if isinstance(key, dpnp_array):
+            idx = key
+        elif isinstance(key, dpt.usm_ndarray):
+            idx = dpnp_array._create_from_usm_ndarray(key)
+        else:
+            try:
+                idx = numpy.asarray(key)
+            except (TypeError, ValueError):
+                return None  # let regular indexing raise
+
+        if dpnp.issubdtype(idx.dtype, dpnp.bool):
+            if idx.ndim > 1:
+                raise IndexError(
+                    "too many indices for flat iterator: flat iterator is "
+                    f"1-dimensional, but {idx.ndim} were indexed"
+                )
+
+            # only a 1-D boolean ndarray mask is valid; reject scalars/lists
+            if idx.ndim == 1 and not isinstance(key, list):
+                # an empty mask selects nothing; otherwise sizes must match
+                if idx.size not in (0, self._size):
+                    raise IndexError(
+                        "boolean index did not match indexed flat iterator "
+                        f"along axis 0; size of axis is {self._size} but size "
+                        f"of corresponding boolean axis is {idx.size}"
+                    )
+                return None
+            raise IndexError("boolean indices for iterators are not supported")
+
+        if not dpnp.issubdtype(idx.dtype, dpnp.integer) or idx.size == 0:
+            return None
+
+        # fancy int indices wrap instead of raising, so bounds-check
+        size = self._size
+        hi, lo = int(idx.max()), int(idx.min())
+        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}")
+        return idx
+
+    def _prepare_key(self, key):
+        # normalize a 1-D iterator index (unwrap a 1-elem tuple; reject None
+        # and longer tuples) and return it with the validated positions or None
+        if isinstance(key, tuple) and len(key) == 1:
+            key = key[0]
+        if key is None or (isinstance(key, tuple) and len(key) > 1):
+            raise IndexError(
+                "only integers, slices (`:`), ellipsis (`...`) and integer "
+                "or boolean arrays are valid indices"
+            )
+        return key, self._validate_key(key)
+
+    def _scalar_pos(self, key):
+        # normalize a scalar flat index (wrap negatives) and bounds-check it
+        pos = key + self._size if key < 0 else key
+        if not 0 <= pos < self._size:
+            raise IndexError(
+                f"index {key} is out of bounds for size {self._size}"
+            )
+        return pos
+
+    def _as_dpnp_array(self, x):
+        # coerce to a dpnp array without overriding a device array's queue
+        if isinstance(x, dpnp_array):
+            return x
+        if isinstance(x, dpt.usm_ndarray):
+            return dpnp_array._create_from_usm_ndarray(x)
+        return dpnp.asarray(
+            x, sycl_queue=self._arr.sycl_queue, usm_type=self._arr.usm_type
+        )
 
     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)
+        key, _ = self._prepare_key(key)
+
+        if isinstance(key, int):
+            # scalar fast path: index directly instead of flattening the array
+            # (a bool key was already rejected in _prepare_key)
+            pos = self._scalar_pos(key)
+            return self._arr[numpy.unravel_index(pos, self._arr.shape)].copy()
+
+        res = dpnp.reshape(self._arr, -1)[key]
+        # basic indexing may alias the source, copy if shares the source buffer
+        # pylint: disable=protected-access
+        src = self._arr.get_array().usm_data._pointer
+        if res.get_array().usm_data._pointer == src:
+            res = res.copy()
+        return res
 
     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)
+        key, idx = self._prepare_key(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; an integer-array key is already
+        # validated (idx), and dpnp.put resolves its negative positions
+        if isinstance(key, int):
+            # scalar fast path: avoid building a full index array
+            # (a bool key was already rejected in _prepare_key)
+            pos = self._scalar_pos(key)
+            idx = dpnp.asarray(pos, sycl_queue=exec_q, usm_type=usm_type)
+        elif isinstance(key, slice):
+            # slice fast path: build only the selected positions
+            start, stop, step = key.indices(self._size)
+            idx = dpnp.arange(
+                start, stop, step, sycl_queue=exec_q, usm_type=usm_type
+            )
+        elif hasattr(key, "dtype") and dpnp.issubdtype(key.dtype, dpnp.bool):
+            # boolean mask fast path
+            mask = self._as_dpnp_array(key)
+            idx = dpnp.nonzero(mask)[0]
+        elif idx is None:
+            # ellipsis, empty tuple or an unrecognized key: let regular
+            # indexing resolve the positions and raise on an invalid key
+            flat_index = dpnp.arange(
+                self._size, sycl_queue=exec_q, usm_type=usm_type
+            )
+            idx = flat_index[key]
+
+        if not dpnp.isscalar(val):
+            val = self._as_dpnp_array(val)
+            if idx.ndim == 0 and val.ndim != 0:
+                # a scalar index targets a single item, reject an array value
+                raise ValueError("Error setting single item of array.")
+
+            val = val.ravel()
+            n = idx.size
+            if val.size not in (0, 1, n):
+                # put broadcasts size-0/1 values; otherwise cycle over 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[self._i]
+            self._i += 1
             return val
         else:
             raise StopIteration