# Batched index\_select

**URL:** <https://discuss.pytorch.org/t/batched-index-select/9115>\
**Category:** Uncategorized\
**Created:** [October 26, 2017, 8:34pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115 "2017-10-26T20:34:41Z")\
**Posts on this page:** 14\
**Page:** 1

<div class="post-metadata">

**Author:** ![Vlad\_Zhukov](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vlad_zhukov/32/4074_2.png) [@Vlad\_Zhukov](https://discuss.pytorch.org/u/Vlad_Zhukov)\
**Post date:** [October 26, 2017, 8:34pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/1 "2017-10-26T20:34:41Z")

</div>

How can I use torch.index\_select for 3D tensors?  
I need some function which will apply usual index\_select per batch.  
F.e. I want to be able to do something like this:

```auto
A = torch.FloatTensor([[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6]], [[7, 8, 9], [7,8, 9]]])
ind = torch.LongTensor([[0, 1], [1, 2], [0, 2]])
torch.index_select(A, 2, ind)

```

What is the most effective and simple way to do this?  
t.y.

---

<div class="post-metadata">

**Author:** ![richard](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/richard/32/1861_2.png) [@richard](https://discuss.pytorch.org/u/richard)\
**Post date:** [October 26, 2017, 11:36pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/2 "2017-10-26T23:36:11Z")

</div>

I’m not sure if this is the best way to do it, but this works:

```auto
torch.cat([torch.index_select(a, 1, i).unsqueeze(0) for a, i in zip(A, ind)])

```

Essentially what this does is apply the regular index\_select to each batch-element of A and ind and then concatenate the results together.

---

<div class="post-metadata">

**Author:** ![Vlad\_Zhukov](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vlad_zhukov/32/4074_2.png) [@Vlad\_Zhukov](https://discuss.pytorch.org/u/Vlad_Zhukov)\
**Post date:** [October 27, 2017, 5:47am UTC](https://discuss.pytorch.org/t/batched-index-select/9115/3 "2017-10-27T05:47:46Z")

</div>

yeah, it’s pretty straightforward. I somehow thought that we can do it without “for loops”…I’m probably overextending 🙂 thank you.

---

<div class="post-metadata">

**Author:** ![wasiahmad](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/wasiahmad/32/248_2.png) [@wasiahmad](https://discuss.pytorch.org/u/wasiahmad)\
**Post date:** [December 5, 2017, 11:22pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/4 "2017-12-05T23:22:22Z")

</div>

Is there any solution of batched `index_select` without using a loop?

---

<div class="post-metadata">

**Author:** ![Vlad\_Zhukov](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vlad_zhukov/32/4074_2.png) [@Vlad\_Zhukov](https://discuss.pytorch.org/u/Vlad_Zhukov)\
**Post date:** [March 25, 2018, 3:26pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/5 "2018-03-25T15:26:53Z")

</div>

This is a bit out of context. But I hope one can figure out what is going here

```auto
def get_selected_matrices(probs, references, dim=1, sanity_check=False):
    """
    batched index select
    probs - is a matrix
    references - is index
    dim - is dimention of element of the batch
    """
    #return torch.cat([torch.index_select(a, dim, Variable(LongTensor(i))).unsqueeze(0)\
    # for a, i in zip(probs, references)])

    batch_size, seq_len, vocab_size = probs.size()
    references = torch.from_numpy(np.array(references)).long() # batch_size x seq_len
    vocab_extension = torch.arange(0, batch_size).long() * vocab_size # batch_size
    if torch.cuda.is_available():
        references = references.cuda()
        vocab_extension = vocab_extension.cuda()
    references_extended_vocab = (references + vocab_extension.unsqueeze(-1)).view(-1) # batch_size * seq_len
    probs_extended_vocab = torch.transpose(probs, 0, 1).contiguous().view(seq_len, -1) # seq_len x batch_size * vocab_size
    probs_reduced_extended_vocab = torch.index_select(
        probs_extended_vocab, dim, Variable(references_extended_vocab)
    ) # seq_len x batch_size * seq_len
    probs_reduced_vocab = torch.transpose(
        probs_reduced_extended_vocab.view(seq_len, batch_size, seq_len), 0, 1
    ) # batch_size x seq_len x seq_len
    if sanity_check:
        probs_reduced_vocab_loop = torch.cat(
            [torch.index_select(a, dim, Variable(LongTensor(i))).unsqueeze(0) for a, i in zip(probs, references)]
        )
        if not torch.equal(probs_reduced_vocab, probs_reduced_vocab_loop):
            raise AssertionError('got wrong probs with reduced vocab')
            print(probs_reduced_vocab)
            print(probs_reduced_vocab_loop)
    return probs_reduced_vocab

```

---

<div class="post-metadata">

**Author:** ![bobchennan](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bobchennan/32/507_2.png) [@bobchennan](https://discuss.pytorch.org/u/bobchennan)\
**Post date:** [April 6, 2018, 5:38pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/6 "2018-04-06T17:38:54Z")

</div>

This is just converting two dimension index into one dimension.

---

<div class="post-metadata">

**Author:** ![jnhwkim](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/jnhwkim/32/1036_2.png) [@jnhwkim](https://discuss.pytorch.org/u/jnhwkim)\
**Post date:** [April 9, 2018, 2:53pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/7 "2018-04-09T14:53:14Z")

</div>

Is this what you want? [`gather`](http://pytorch.org/docs/master/torch.html#torch.gather) might be useful for the task

```auto
# Batched index_select
def batched_index_select(t, dim, inds):
    dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2))
    out = t.gather(dim, dummy) # b x e x f
    return out

```

---

<div class="post-metadata">

**Author:** ![marcman411](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/7cd45c/32.png) [@marcman411](https://discuss.pytorch.org/u/marcman411)\
**Post date:** [May 10, 2018, 6:25pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/8 "2018-05-10T18:25:38Z")

</div>

Piggy-backing on @jnhwkim’s idea of using gather, here is a function that should mimic `index_select` for arbitrary dimensional inputs. It just doesn’t work with negative indexing.

```auto
# input: B x * x ... x *
# dim: 0 <= scalar
# index: M
def batched_index_select(input, dim, index):
	views = [1 if i != dim else -1 for i in xrange(len(input.shape))]
	expanse = list(input.shape)
	expanse[dim] = -1
	index = index.view(views).expand(expanse)
	return torch.gather(input, dim, index)

```

And here’s a variation if you have different indexes for each batch:

```auto
# input: B x * x ... x *
# dim: 0 < scalar
# index: B x M
def batched_index_select(input, dim, index):
	views = [input.shape[0]] + \
		[1 if i != dim else -1 for i in xrange(1, len(input.shape))]
	expanse = list(input.shape)
	expanse[0] = -1
	expanse[dim] = -1
	index = index.view(views).expand(expanse)
	return torch.gather(input, dim, index)

```

---

<div class="post-metadata">

**Author:** ![ratishsp](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/e480ec/32.png) [@ratishsp](https://discuss.pytorch.org/u/ratishsp)\
**Post date:** [July 18, 2018, 3:48pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/9 "2018-07-18T15:48:45Z")

</div>

> [@jnhwkim](#):
>
> def batched\_index\_select(t, dim, inds): dummy = inds.unsqueeze(2).expand(inds.size(0), inds.size(1), t.size(2)) out = t.gather(dim, dummy) # b x e x f return out

Hi @jnhwkim,  
The expected output is  
(0 ,.,.) =  
1 2  
1 2

(1 ,.,.) =  
5 6  
5 6

(2 ,.,.) =  
7 9  
7 9  
[torch.FloatTensor of size 3x2x2]

Running your version of batched\_index\_select results in a different output  
(0 ,.,.) =  
1 1 1  
2 2 2

(1 ,.,.) =  
5 5 5  
6 6 6

(2 ,.,.) =  
7 7 7  
9 9 9  
[torch.FloatTensor of size 3x2x3]

---

<div class="post-metadata">

**Author:** ![Zhaoyi-Yan](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zhaoyi-yan/32/8480_2.png) [@Zhaoyi-Yan](https://discuss.pytorch.org/u/Zhaoyi-Yan)\
**Post date:** [July 23, 2019, 8:33am UTC](https://discuss.pytorch.org/t/batched-index-select/9115/10 "2019-07-23T08:33:12Z")

</div>

Maybe a slightly better version is here.

```auto
import torch
# input: B x * x ... x *
# dim: 0 <= scalar
# index: M
def batched_index_select(input, dim, index):
    views = [1 if i != dim else -1 for i in range(len(input.shape))]
    expanse = list(input.shape)
    expanse[dim] = -1
    index = index.view(views).expand(expanse)
    # making the first dim of output be B
    return torch.cat(torch.chunk(torch.gather(input, dim, index), chunks=index.shape[0], dim=dim), dim=0)

```

---

<div class="post-metadata">

**Author:** ![dashesy](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dashesy/32/4869_2.png) [@dashesy](https://discuss.pytorch.org/u/dashesy)\
**Post date:** [September 10, 2019, 1:13am UTC](https://discuss.pytorch.org/t/batched-index-select/9115/11 "2019-09-10T01:13:34Z")

</div>

> [@marcman411](#):
>
> different indexes for each batch

That is exactly what I wanted, thanks. Since I did not want to rely on the exact batch number I used this variant

```python
def batched_index_select(input, dim, index):
    for ii in range(1, len(input.shape)):
        if ii != dim:
            index = index.unsqueeze(ii)
    expanse = list(input.shape)
    expanse[0] = -1
    expanse[dim] = -1
    index = index.expand(expanse)
    return torch.gather(input, dim, index)

```

---

<div class="post-metadata">

**Author:** ![spadel](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/8797f3/32.png) [@spadel](https://discuss.pytorch.org/u/spadel)\
**Post date:** [January 20, 2021, 7:04pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/12 "2021-01-20T19:04:21Z")

</div>

I have a similar task, so maybe someone can help me!

So what I have are two tensors:

- an indices tensor `indices` with shape `(2, 5, 2)`, where the last dimensions corresponds to indices in x and y dimension
- a “value tensor” `value` with shape `(2, 5, 2, 16, 16)`, where I want the last two dimensions to be selected with x and y indices

To be more concrete, the indices are between 0 and 15 and I want to get an output:

```
out = value[:, :, :, x_indices, y_indices]

```

For example the indices for the first item are `i_y = indices[0, 0, 0]; i_x = indices[0, 0, 1]`. The corresponding output should then be `out[0, 0, :, i_y, i_x]`. Can anybody help me here? Thanks a lot!

---

<div class="post-metadata">

**Author:** ![eacousineau](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/eacousineau/32/26388_2.png) [@eacousineau](https://discuss.pytorch.org/u/eacousineau)\
**Post date:** [December 3, 2021, 6:14pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/13 "2021-12-03T18:14:38Z")

</div>

Not sure if this is comparable to [dashesy](https://discuss.pytorch.org/u/dashesy)’s in terms of generality, but here’s a simplified case + tests for selecting vectors, e.g. via `argmin` or `multinomial`:

> <https://gist.github.com/EricCousineau-TRI/cc2dc27c7413ea8e5b4fd9675050b1c0>

---

<div class="post-metadata">

**Author:** ![Zohaib](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/z/3be4f8/32.png) [@Zohaib](https://discuss.pytorch.org/u/Zohaib)\
**Post date:** [February 8, 2022, 11:09pm UTC](https://discuss.pytorch.org/t/batched-index-select/9115/14 "2022-02-08T23:09:47Z")

</div>

Thanks, @eacousineau. I was looking for this.