Negative index support for `take_along_dim`

[mdhaber]

🚀 The feature, motivation and pitch

I'm working on adding an implementation of quantile in terms of Python array API standard calls¹ for SciPy (scipy/scipy#22352), and I would like use of negative indices to be possible in torch.take_along_dim.

import torch as xp
x = xp.asarray([1, 2, 3])
xp.take(x, xp.asarray(-1))  # tensor(3)
xp.take_along_dim(x, xp.asarray([-1]))  # expected tensor(3), but got
# RuntimeError: index -1 is out of bounds for dimension 0 with size 3

On the GPU, we get:

import torch as xp
device = "cuda" if xp.cuda.is_available() else "cpu"
x = xp.asarray([1, 2, 3], device=device)
xp.take_along_dim(x, xp.asarray(-1, device=device))
# RuntimeError: CUDA error: device-side assert triggered
# CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
# For debugging consider passing CUDA_LAUNCH_BLOCKING=1
# Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Alternatives

array_api_compat can patch this, or I can always calculate the equivalent positive index.

Additional context

No response

cc @mruberry @rgommers @asmeurer @leofang @AnirudhDagar @asi1024 @emcastillo @kmaehashi @albanD

Footnotes

1. take_along_axis will be available in the next version of the standard. It is not explicit in the take_along_axis documentation about negative indices, but negative indices seem to be supported in general. In any case, I've asked for clarification. ↩

[albanD]

Hey!
Would adding % x.size(0) as xp.take_along_dim(x, xp.asarray([-1]) % x.size(0)) work for you?
But if the array API specs says that we should handle negative indices in there, then happy to review a PR adding this feature!

[mdhaber]

I address the possible workarounds in "Alternatives" above, and I've already gotten past it in my use case. But it looks like Torch is working toward array API compliance, so I thought it should be reported. Thanks for considering it!

Update: confirmed that intent of standard is for negative index support throughout.

[albanD]

Thanks for the link. Would definitely accept a PR adding this!

[rgommers]

I think dim accepting negative values is also generally expected. Other functions do that, e.g.:

>>> t = torch.tensor([[0, 1, 2], [0, 1, 2]])
>>> torch.unique(t, dim=0)
tensor([[0, 1, 2]])
>>> torch.unique(t, dim=1)
tensor([[0, 1, 2],
        [0, 1, 2]])
>>> torch.unique(t, dim=-1)  # `dim=-1` specifies last axis, often clearer than a positive value
tensor([[0, 1, 2],
        [0, 1, 2]])

There's some dim normalization utility function that take_along_axis is failing to call I suspect.

[KaaustaaubShankar]

@albanD Hey! I put a pr up for this issue so if you could take a look, I would appreciate it!