Skip to content

Conversation

kgryte
Copy link
Contributor

@kgryte kgryte commented Feb 6, 2025

This PR

  • clarifies that negative indices should be supported in take and take_along_axis. This aligns function behavior with integer indexing, as currently specified in the specification.
  • rearranges content to move notes to a dedicated notes section.

Support for negative indices is supported in NumPy et al. E.g., in NumPy

In [1]: import numpy as np

In [2]: x = np.ones(10,dtype="float32")

In [3]: np.take_along_axis(x, np.asarray([-1]), axis=0)
Out[3]: array([1.], dtype=float32)

In [4]: np.take_along_axis(x, np.asarray([-2]), axis=0)
Out[4]: array([1.], dtype=float32)

In [5]: np.take_along_axis(x, np.asarray([-3]), axis=0)
Out[5]: array([1.], dtype=float32)

In [6]: np.take_along_axis(x, np.asarray([-10]), axis=0)
Out[6]: array([1.], dtype=float32)

In [7]: np.take_along_axis(x, np.asarray([-11]), axis=0)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[7], line 1
----> 1 np.take_along_axis(x, np.asarray([-11]), axis=0)

File ~/anaconda3/lib/python3.10/site-packages/numpy/lib/_shape_base_impl.py:173, in take_along_axis(arr, indices, axis)
    170     arr_shape = arr.shape
    172 # use the fancy index
--> 173 return arr[_make_along_axis_idx(arr_shape, indices, axis)]

IndexError: index -11 is out of bounds for axis 0 with size 10

PyTorch does not (yet) support negative indices, but is willing to take a PR adding support: pytorch/pytorch#146211.

cc @mdhaber

@kgryte kgryte added Maintenance Bug fix, typo fix, or general maintenance. topic: Indexing Array indexing. labels Feb 6, 2025
@kgryte kgryte added this to the v2024 milestone Feb 6, 2025
Copy link
Contributor

@mdhaber mdhaber left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@kgryte
Copy link
Contributor Author

kgryte commented Feb 6, 2025

Thanks for the review, @mdhaber! In it goes!

@kgryte kgryte merged commit df8081c into data-apis:main Feb 6, 2025
3 checks passed
@kgryte kgryte deleted the fix/take_along_axis_negative_indices branch February 6, 2025 05:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Maintenance Bug fix, typo fix, or general maintenance. topic: Indexing Array indexing.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants