Fix dpnp.ndarray.flat indexing edge cases

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
 

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,180 @@
 
 """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.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 _unwrap_tuple(key):
+        # a flat iterator is 1-D, so a single-element index tuple is equivalent
+        # to its element (e.g. `flat[(idx,)]` behaves like `flat[idx]`)
+        if isinstance(key, tuple) and len(key) == 1:
+            return key[0]
+        return key
+
+    @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
+
+        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 Exception:
+                return  # let regular indexing raise
+
+        if not dpnp.issubdtype(idx.dtype, dpnp.integer) or idx.size == 0:
+            return
+
+        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}")
+
+    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)
+        key = self._unwrap_tuple(key)
+        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)
+        key = self._unwrap_tuple(key)
+        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
+        if isinstance(key, int) and not isinstance(key, bool):
+            # fast path for a scalar index: avoid building a full index array
+            pos = key + a.size if key < 0 else key
+            if not 0 <= pos < a.size:
+                raise IndexError(
+                    f"index {key} is out of bounds for size {a.size}"
+                )
+            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(a.size)
+            idx = dpnp.arange(
+                start, stop, step, sycl_queue=exec_q, usm_type=usm_type
+            )
+        else:
+            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