Fix dpnp.ndarray.flat indexing edge cases - #3045
Conversation
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).
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3045/index.html |
|
Array API standard conformance tests for dpnp=0.21.0dev7=py314ha0e2e8e_16 ran successfully. |
|
coverage: 78.5% (+0.02%) from 78.483% — fix/SAT-8204-flatiter-indexing into master |
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.
|
|
||
| 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)): |
There was a problem hiding this comment.
This logic looks wrong for a tuple
If we pass a tuple index OOB index will not be caught
In [1]: import numpy, dpnp
In [2]: a = numpy.array([1,2,3])
In [3]: a_dp = dpnp.array(a)
In [4]: a.flat[(numpy.array([5]),)]
---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Cell In[4], line 1
----> 1 a.flat[(numpy.array([5]),)]
IndexError: index 5 is out of bounds for size 3
In [5]: a_dp.flat[(dpnp.array([5]),)]
Out[5]: array([3])
There was a problem hiding this comment.
The same as for bool and gives a different shape
In [6]: a.flat[True]
<ipython-input-6-e68040906b52>:1: DeprecationWarning: Indexing flat iterators with a 0-dimensional boolean index is deprecated and may be removed in a future version. (Deprecated NumPy 2.4)
a.flat[True]
Out[6]: np.int64(1)
In [7]: a_dp.flat[True]
Out[7]: array([[1, 2, 3]])
There was a problem hiding this comment.
Regarding flat[True] (the 0-d boolean case): numpy deprecated it in 2.4 (DeprecationWarning: … may be removed in a future version), so it seemed low-value to add code replicating a quirk numpy is actively removing.
There was a problem hiding this comment.
So I'd prefer to keep the standard array boolean-index semantics, rather than duplicating the deprecated numpy behavior.
Once the deprecation expires, numpy will also return:
a.flat[True]
# Out: array([[1, 2, 3]])There was a problem hiding this comment.
The tuple index OOB is addressed.
| return | ||
|
|
||
| size = self._size | ||
| hi, lo = int(dpnp.max(idx)), int(dpnp.min(idx)) |
There was a problem hiding this comment.
2 copy to host and .wait() executions
Is there any way to reduce the host transfer here?
There was a problem hiding this comment.
Only left for the use case when input key is dpnp.ndarray or usm.ndarray.
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.
3 participants
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.