Merge pull request #361 from ev-br/torch_take_along_axis_neg_idx

ENH: torch: allow negative indices in `take` and `take_along_axis`

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -815,11 +815,24 @@ def take(x: Array, indices: Array, /, *, axis: int | None = None, **kwargs: obje
         if x.ndim != 1:
             raise ValueError("axis must be specified when ndim > 1")
         axis = 0
-    return torch.index_select(x, axis, indices, **kwargs)
+    # torch does not support negative indices,
+    # see https://github.com/pytorch/pytorch/issues/146211
+    return torch.index_select(
+        x,
+        axis,
+        torch.where(indices < 0, indices + x.shape[axis], indices),
+        **kwargs
+    )
 
 
 def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
-    return torch.take_along_dim(x, indices, dim=axis)
+    # torch does not support negative indices,
+    # see https://github.com/pytorch/pytorch/issues/146211
+    return torch.take_along_dim(
+        x,
+        torch.where(indices < 0, indices + x.shape[axis], indices),
+        dim=axis
+    )
 
 
 def sign(x: Array, /) -> Array: