# Is there any better way to do this kind of indexing?

**URL:** <https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213>\
**Category:** Uncategorized\
**Created:** [February 26, 2020, 10:42pm UTC](https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213 "2020-02-26T22:42:16Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![entslscheia](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/entslscheia/32/6959_2.png) [@entslscheia](https://discuss.pytorch.org/u/entslscheia)\
**Post date:** [February 26, 2020, 10:42pm UTC](https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213/1 "2020-02-26T22:42:16Z")

</div>

For example, I have a tensor `a` of shape `[k, m, n]` and tensor `b` fo shape `[k,]`. What I want to do is to get `a[torch.arange(k), b]`. Is there any api provided by PyTorch can directly deal with it? Creating a new tensor using `torch.arange(k)` doesn’t look that elegant.

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [February 26, 2020, 10:54pm UTC](https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213/2 "2020-02-26T22:54:00Z")

</div>

Hi,

I’m confused what you expect to get here. what does b contain?  
Can you share a small code sample that shows what you expect as output?

---

<div class="post-metadata">

**Author:** ![entslscheia](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/entslscheia/32/6959_2.png) [@entslscheia](https://discuss.pytorch.org/u/entslscheia)\
**Post date:** [February 26, 2020, 11:02pm UTC](https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213/3 "2020-02-26T23:02:42Z")

</div>

Sorry, I should have included more information here. Every element in `b` is in `[0, m)`. You can imagine that we have a sequence ranking task. k is the batch size, and for each batch instance, we have m candidate sequences of length n. So now `b` is the index of candidates, and there are k of them, because we want to choose one candidate for every batch instance. I think there should be some more straightforward way to handle this right? Since it kinda seems like a common need.

---

<div class="post-metadata">

**Author:** ![albanD](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/32/215_2.png) [@albanD](https://discuss.pytorch.org/u/albanD)\
**Post date:** [February 26, 2020, 11:11pm UTC](https://discuss.pytorch.org/t/is-there-any-better-way-to-do-this-kind-of-indexing/71213/4 "2020-02-26T23:11:28Z")

</div>

Yes, you can do: `a.gather(1, b.view(k, 1, 1).expand(k, 1, n))` will give you what you want 🙂  
Basically, use the gather function, and since you want all the elements over the last dimension of size n, you just expand the Tensor of indices there.