pythonindexingpytorchtensortorch

How to map values from a 3D tensor to a 1D tensor in PyTorch?

I'm stuck with a Pytorch problem and could use some help:

I've got two tensors:

I need to create a new tensor that's the same shape as the first one (i, j, j), but where each value is replaced by the corresponding value from the second tensor. Like using the values in the first tensor as indices for the second one.

Here's a quick example:

import torch

# My tensors
big_tensor = torch.randint(0, 256, (10, 25, 25))
small_tensor = torch.rand(256)

# What I'm trying to do
result = magic_function(big_tensor, small_tensor)

# How it should work
print(big_tensor[0, 0, 0])    # Let's say this outputs 42
print(small_tensor[42])       # This might output 0.7853
print(result[0, 0, 0])        # This should also be 0.7853

I'm looking for a performant way to do this, preferably not using loops, as both tensors can be quite big. Is there an efficient Pytorch operation or method I can use for this?

Any help would be awesome - thanks in advance!

Solution

  • This should work:

    small_tensor[big_tensor]

    Take note that the type of the big_tensor must be long/int.

    Edit:

    In response to the comment of @simon, I wrote a colab notebook that shows how this solution works without the need to perform any other operation.