Fix dpnp.ndarray.flat indexing edge cases - #3045
Open
antonwolfy wants to merge 41 commits into
Open
Conversation
Contributor
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3045/index.html |
Contributor
|
Array API standard conformance tests for dpnp=0.21.0dev8=py314ha0e2e8e_49 ran successfully. |
Collaborator
|
coverage: 78.541% (+0.06%) from 78.477% — fix/SAT-8204-flatiter-indexing into master |
antonwolfy
marked this pull request as ready for review
August 28, 2026 16:37
antonwolfy
requested review from
ndgrigorian and
vlad-perevezentsev
as code owners
August 28, 2026 16:37
ndgrigorian
reviewed
Sep 1, 2026
Support slices, ellipsis, empty tuple, and integer/boolean array indices in flatiter __getitem__/__setitem__, reusing regular array indexing for validation. Reject numpy.newaxis (None), raise IndexError for out-of-bounds integer array indices, and reject assignment with a 0-D index, matching NumPy (gh-28590). Return copies from __getitem__ rather than views. SAT-8204
Cover slices, ellipsis, empty tuple, integer/boolean array indices, newaxis rejection, out-of-bounds, non-contiguous write-back, and copy-not-view semantics. Cross-check against NumPy where behavior is shared, and gate NumPy 2.4-only cases (numpy-gh-28590) with testing.with_requires.
Match numpy's np.put-style cycling when a flat assignment value is shorter than the selection (dpnp.put broadcasts instead). Enable the previously-disabled slice/ellipsis/empty-tuple parametrizations in the cupy flatiter iterate tests, and drop the IndexError cases that became valid indices in numpy 2.4 (numpy-gh-28590).
Expand the flatiter and ndarray.flat docstrings to align with NumPy, documenting supported basic and advanced indexing and adding See Also and Examples sections. Fix a broken dpnp.flat cross-reference in the ndarray.flatten docstring.
Resolve a scalar integer or slice flat index to positions directly instead of allocating arange(size) and indexing it, so a single-element or slice assignment no longer materializes a full index array.
A 1-D flat iterator takes a single index, so unwrap a 1-element index tuple to its element before validation. This makes an out-of-bounds array index wrapped in a tuple (e.g. arr.flat[(array([5]),)]) raise IndexError as in NumPy, instead of silently wrapping.
Validate an out-of-bounds flat index by inspecting the raw index on the host (numpy) when it is not already a device array, instead of always uploading it via dpnp.asarray and reducing on device. Add tests for usm_ndarray and empty index keys.
A scalar flat index targets a single element, so assigning an array (ndim >= 1) value now raises ValueError to match NumPy and dpnp's own scalar element assignment, instead of silently taking the first value.
Gate the single-item array-assignment test on numpy>=2.4 (the 0-d array index only raises there), and remove the numpy>=2.4 gate from the boolean-mask and array-out-of-bounds tests, which already behave identically on older NumPy.
A flat iterator is 1-D and takes a single index, so a tuple with more than one element (e.g. arr.flat[..., 2]) now raises IndexError as in NumPy, instead of silently absorbing the extra dimensions and, for setitem, mutating the array.
Catch only TypeError/ValueError from numpy.asarray when resolving a flat index for bounds-checking, so unexpected errors propagate. Add a ragged-index test covering the fallback.
Fold tuple-unwrap and invalid-key rejection into _normalize_key (called by getitem/setitem), inline the single-use flat view, and abbreviate a comment.
A boolean flat index is only valid as an ndarray mask; a boolean scalar (True/False or 0-d bool array) or a boolean list now raises IndexError, matching NumPy (gh-28590) and avoiding a whole-array overwrite on assignment. Re-enable the boolean-scalar case in the cupy iterate IndexError tests and add dpnp tests for the rejected forms.
Gate the single-element-tuple and tuple-wrapped out-of-bounds tests on numpy>=2.4, since tuple indexing of a flat iterator became valid only in NumPy 2.4 (numpy-gh-28590). Assert the iterated element count so an early-stop __next__ regression is caught despite zip truncation.
antonwolfy
force-pushed
the
fix/SAT-8204-flatiter-indexing
branch
from
September 3, 2026 15:41
f7c9349 to
e314321
Compare
Single-element tuple indexing and tuple-wrapped out-of-bounds already behave the same on NumPy 2.0, so drop the unnecessary numpy>=2.4 gate to keep the tests running on the older-NumPy CI leg.
- __getitem__: add a scalar fast path that resolves the multi-index and indexes the array directly, avoiding a reshape (which copies a non-contiguous array) per element during iteration. - _validate_key: drop the isinstance(int) clause subsumed by the __index__ check, and simplify the boolean guard (bool/tuple were unreachable); reject a boolean mask whose size does not match. - __setitem__: resolve a boolean mask via dpnp.nonzero instead of arange(size)[mask], avoiding the full-size index buffer. - tests: cover negative and non-contiguous scalar get, multiple bool mask input types, and a wrong-size bool mask.
NumPy only rejects a size-mismatched boolean flat index since 2.4 (numpy-gh-28590); on older NumPy it treats the mask as integer indices, so gate the test to avoid a failure on the numpy 2.0 CI leg.
An empty boolean flat index selects nothing (get returns an empty array, set is a no-op) in NumPy, so exempt a zero-size mask from the size-match check instead of raising. Add a test.
Dedupe the scalar index normalization/bounds-check shared by __getitem__ and __setitem__ into _scalar_pos, replace a chained comparison with a clearer form, and rename a boolean-mask variable.
The scalar-int early return is bounds-checked later by _scalar_pos, not by regular indexing.
dpnp validates a fancy index even when the assigned value is empty, raising IndexError where NumPy no-ops; add a dpnp-only test to lock in this intentional divergence.
| % val.size | ||
| ] | ||
|
|
||
| dpnp.put(a, idx, val) |
Contributor
There was a problem hiding this comment.
Calling dpnp.put prevents memory overlap while NumPy allows it
In [10]: a = dpnp.arange(6)
In [11]: a_np = numpy.arange(6)
In [12]: a.flat[:-1] = a[1:]
ValueError: Arrays index overlapping segments of memory
In [13]: a_np.flat[:-1] = a_np[1:]
In [14]: a[:-1] = a[1:] # setitem works
Contributor
Author
There was a problem hiding this comment.
dpnp.ndarray.flat setitem is built on dpnp.put, which rejects overlapping memory by design, so this is consistent with the rest of dpnp rather than a flat-specific gap.
Matching NumPy here isn't cleanly feasible either: NumPy's flat setitem cascades on a backward overlap (a.flat[1:] = a[:-1] -> [0,0,0,0,0,0]), which a gather/scatter put can't reproduce; the only workaround (copying the value) yields snapshot semantics that match NumPy's regular assignment but not its flat behavior.
Given that, I'd prefer to keep raising a clear error rather than add a masked copy that still diverges.
A flat iterator is 1-D, so a boolean mask with ndim > 1 is too many indices and now raises IndexError as in NumPy, instead of silently writing wrong positions (the setitem path used nonzero(mask)[0], i.e. only the first-axis indices). Add a test.
Advanced indexing already allocates a fresh array, so an unconditional .copy() doubled the allocation. Copy only when the result aliases the source (shares the base allocation `_pointer`), i.e. a basic-index view. Extend the copy-semantics test to cover a non-contiguous source.
The previous commit compared usm_ndarray._pointer, which is adjusted to the first element, so a slice view (non-zero offset) compared unequal to its source and the copy was wrongly skipped, leaking a writable view. Compare usm_data._pointer (the offset-independent allocation base) instead, matching what "same underlying buffer" actually means.
Passing sycl_queue/usm_type to dpnp.asarray() silently migrated a value or boolean mask that lived on another queue, so flat setitem accepted cross-queue inputs that regular assignment, dpnp.put and flat getitem all reject. Add a _asarray_cfd helper that keeps a device array on its own queue (letting the downstream put/nonzero raise ExecutionPlacementError) and only places a host input on the iterator's queue.
Rename the setitem coercion helper to _as_dpnp_array and return a dpnp_array unchanged (wrapping a usm_ndarray without a copy), trimming its comment. Move the cross-queue negative test out of test_flat into test_sycl_queue as test_flat_setitem_diff_queue, and add positive persistence tests: test_flat in test_sycl_queue (queue stays with the input across flat get/setitem) and test_flat in test_usm_type (slice keeps the usm type, advanced index coerces, setitem preserves it).
dpnp.put broadcasts size-0/1 (and 0-d) values across the target indices, matching NumPy's flatiter assignment which cycles a value iterator that resets when exhausted (a 1-element value fills the whole selection). So cycling a size-1 value built an arange(n) + modulo + gather for a result put already produces. Cycle only for a genuine mismatch (val.size not in (0, 1, n)) and add a test pinning the size-1/0-d broadcast to NumPy.
_normalize_key already raises for any boolean key, so key is never a bool at the scalar fast path; the extra `not isinstance(key, bool)` (guarding against bool being an int subclass) can never fire. Replace it with `isinstance(key, int)` in getitem and setitem, noting why.
_validate_key already materializes and bounds-checks the index array for an integer-array key, so return it (renaming _normalize_key to _prepare_key, which now yields both the key and its positions). The setitem integer-array path passes that idx straight to dpnp.put, whose mode="wrap" resolves the in-range negatives, instead of building a size-sized arange and gathering into it; arange is now only the fallback for ellipsis or an unrecognized key. Add a test covering an integer-array index with a per-element value for both positive and negative positions.
NumPy says the boolean index did not match the "indexed flat iterator" along axis 0; dpnp said "indexed array". Align the message and assert it in test_flat_bool_mask_wrong_size.
The flat entry covered only indexing; the branch also fixes assignment (compute-follows-data, boolean-mask writes, value cycling/broadcast, and the 0-D-index rejection).
Cover cases the suite missed: value cycling over a longer selection (n>1), an in-bounds empty selection with an empty value, a non-integer array index (the validate return-None path), setitem via a 1-element tuple, and the slice-setitem queue-placement path in test_sycl_queue.
The slice and fallback paths read a.size while the rest of the class uses the cached self._size; standardize on self._size.
Use self[self._i] and self._i += 1 instead of the explicit __getitem__ call and rebind.
Contributor
Author
|
Ready for the next review round. Since the last one:
All flatiter tests pass locally. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR reworks
dpnp.flatiter(dpnp.ndarray.flat) indexing so it aligns with NumPy's flat-iterator semantics, which were tightened in NumPy 2.4.Previously
dpnp.ndarray.flataccepted only a single integer index and raisedTypeErrorfor everything else. Indexing now delegates to regular array indexing of the flattened array, so it supports the full set of flat index types and matches NumPy's error behavior.