Negative index support for take_along_dim #146211
Description
🚀 The feature, motivation and pitch
I'm working on adding an implementation of quantile in terms of Python array API standard calls1 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 3On 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
-
take_along_axiswill be available in the next version of the standard. It is not explicit in thetake_along_axisdocumentation about negative indices, but negative indices seem to be supported in general. In any case, I've asked for clarification. ↩