Understanding torch.gather function in Pytorch

Pranav Chaturvedi
Oct 18 · 3 min read

Two arguments of this function, index and dim are the key to understanding the function.

For case of 2D, dim = 0 corresponds to rows and dim = 1 corresponds to columns.

For case of 3D, dim = 0 corresponds to image from batch, dim = 1 corresponds to rows and dim = 2 corresponds to columns.

Case of 2D input tensor

1. Understanding dim argument:

a. When dim = 0, we choose rows.

b. When dim = 1, we choose columns.

2. Understanding index argument:

a. Index argument will have same no of dimensions as input(does not mean shape will be same).

b. Output tensor will have the same shape as index tensor.

c. The elements of index tensor tell which row (for dim = 0, 2D case) to choose and position of the particular element tells which column to choose.

d. The elements of index tensor tell which column (for dim = 1, 2D case) to choose and position of the particular element tells which row to choose.

Case of 3D input tensor

1. Understanding dim argument:

a. When dim = 0, we choose image from batch.

b. When dim = 1, we choose rows.

c. When dim = 2, we choose columns.

2. Understanding index argument:

a. a and b from above, Case of 2D input tensor apply.

b. The elements of index tensor tell which image from batch (for dim = 0, 3D case) to choose and position of the particular element tells which rows and columns to choose and so on for dim = 1 and dim = 2.

Let’s take two examples for case of 2D.

1st Example

When dim = 0, and

ind_2d = [[3, 2, 0, 1]]

ind_2d has shape (1, 4) so output will have same shape.

0th element of ind_2d, i.e. 3 tells we choose 3rd row and 0th column (since 3 is 0th element of index tensor).

1st element of ind_2d i.e. 2 tells we choose 2nd row starting from 0 (row because dim = 0) and 1st column (since 1st element of index tensor). And so on.

2nd Example

Let’s suppose from above src_2d tensor we want to select 0, 6, 10 and 15. (Tip: We read these number order-wise from up to down, so we form a column like index tensor. When we read from left to right, we form row like index tensor.)

Now 0 belongs to 0th column, 6 belongs to 2nd column, 10 belongs to 2nd column and 15 belongs to 3rd column.

So our index tensor is [[0, 2, 2, 3]] in column form i.e. of shape (4, 1). And since we’ve selected columns, therefore dim = 1.

There’s no way we could have selected rows and got the desired output tensor.

The case of 3D is very similar. Get the ideas from infographic below.

Image for post
Image for post
Example for the case of 3D input tensor

Here is the link to a python notebook with several examples for both cases of 2D and 3D input tensor:

Link to the Pytorch function :

https://pytorch.org/docs/stable/generated/torch.gather.html?highlight=torch%20gather#torch.gather

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data…

Sign up for Analytics Vidhya News Bytes

By Analytics Vidhya

Latest news from Analytics Vidhya on our Hackathons and some of our best articles! Take a look

By signing up, you will create a Medium account if you don’t already have one. Review our Privacy Policy for more information about our privacy practices.

Check your inbox
Medium sent you an email at to complete your subscription.

Pranav Chaturvedi

Written by

Artificial Intelligence and Machine Learning

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

Pranav Chaturvedi

Written by

Artificial Intelligence and Machine Learning

Analytics Vidhya

Analytics Vidhya is a community of Analytics and Data Science professionals. We are building the next-gen data science ecosystem https://www.analyticsvidhya.com

Medium is an open platform where 170 million readers come to find insightful and dynamic thinking. Here, expert and undiscovered voices alike dive into the heart of any topic and bring new ideas to the surface. Learn more

Follow the writers, publications, and topics that matter to you, and you’ll see them on your homepage and in your inbox. Explore

If you have a story to tell, knowledge to share, or a perspective to offer — welcome home. It’s easy and free to post your thinking on any topic. Write on Medium

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store