Advanced Selection from Tensors in Pytorch

Using torch.index_select, torch.gather and torch.take

In some situations, you’ll need to do some advanced indexing / selection with Pytorch, e.g. answer the question: “how can I select elements from Tensor A following the indices specified in Tensor B?”

In this post we’ll present the three most common methods for such tasks, namely torch.index_select, torch.gather and torch.take. We’ll explain all of them in detail and contrast them with one another.