Efficiently Selecting Values from Tensors in PyTorch: Using Indices from Another Tensor

• You have two PyTorch tensors:
  • a: A tensor with multiple dimensions, but we're particularly interested in the last dimension (often representing features).
  • b: A tensor with a smaller number of dimensions (usually one less than a). This tensor contains indices that will be used to select specific values from the last dimension of a.

Goal:

You want to create a new tensor c that contains the values from a selected based on the indices in b. The resulting tensor c will have the same shape as b.

Methods:

Here are two common approaches to achieve this:

Method 1: Using a loop

1. Iterate over each element in b:

   • Use a loop (e.g., for) to iterate through each element or row in b.

2. Select the corresponding value(s) from a:

3. Store the selected value(s) in c:

Example:

import torch

a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)  # Example tensor a
b = torch.tensor([1, 0])  # Example tensor b (indices)

c = torch.zeros_like(b)  # Create c with the same shape as b

for i in range(b.size(0)):  # Loop through each element in b
    c[i] = a[i, b[i]]  # Select value from a using index from b

print(c)

Method 2: Using gather (more efficient for large tensors)

1. Reshape b (optional):

2. Apply torch.gather:

import torch

a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)  # Example tensor a
b = torch.tensor([1, 0])  # Example tensor b (indices)

c = torch.gather(a, dim=-1, index=b.unsqueeze(-1))  # Gather using b as indices

print(c)

Key Points:

• The last dimension of a typically represents features you want to select from.
• b provides the indices that determine which values to extract from a.
• Both methods achieve the same goal, but torch.gather is generally more efficient for larger tensors.
• Choose the method that best suits your specific needs and performance considerations.

import torch

# Example tensor a with multiple dimensions (here, 3 dimensions)
a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)  
# Example tensor b with indices (smaller dimension)
b = torch.tensor([1, 0]) 

# Create an empty tensor c with the same shape as b to store the selected values
c = torch.zeros_like(b)

# Loop through each element (row) in b
for i in range(b.size(0)):
  # Reshape b temporarily if necessary (ensures correct indexing)
  index = b[i].unsqueeze(-1)  # Add a new dimension at the end

  # Select the corresponding value(s) from the last dimension of a using the index
  c[i] = a[i, index] 

print(c)

Explanation:

1. We import the torch library for PyTorch operations.
2. We create two example tensors:
   • a: A 3D tensor with shape (2, 2, 3), representing some data.
   • b: A 1D tensor with shape (2,), containing indices (here, 1 and 0) to select values from a.
3. We create an empty tensor c with the same shape as b to store the selected values.
4. We iterate through each element (row) in b using a for loop.
5. Inside the loop, we temporarily reshape b[i] by adding a new dimension at the end using unsqueeze(-1). This ensures compatibility with the indexing requirements of a.
6. We use the reshaped index b[i] to access the corresponding value(s) from the last dimension of a using a[i, index].
7. We store the selected value(s) in the corresponding position of c.
8. Finally, we print the resulting tensor c.

import torch

# Example tensor a with multiple dimensions (here, 3 dimensions)
a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)  
# Example tensor b with indices (smaller dimension)
b = torch.tensor([1, 0]) 

# Gather elements from a using b as indices along the last dimension
c = torch.gather(a, dim=-1, index=b.unsqueeze(-1))

print(c)

1. We import the torch library.
2. We create the same example tensors a and b as in Method 1.
3. We use the torch.gather function. It takes three arguments:
   • a: The tensor from which to select elements.
   • dim: The dimension along which to gather (here, the last dimension, specified as -1).
   • index: The tensor containing the indices to select (here, b). We reshape b by adding a new dimension at the end using unsqueeze(-1) for compatibility with gather.
4. The torch.gather function efficiently selects the elements from the last dimension of a based on the corresponding indices in b.
5. The result is stored in the new tensor c.

• torch.index_select can be used for a specific scenario where you want to select entire rows or columns based on indices in b. However, it's not directly applicable for selecting individual elements within the last dimension.

Example (limited to selecting rows/columns):

import torch

# Example tensors (assuming b has the same number of dimensions as a except last)
a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)
b = torch.tensor([1, 0])

# Select rows from a based on indices in b (requires same number of dimensions)
c = torch.index_select(a, dim=0, index=b)

print(c)

Using List Comprehension (Less Efficient):

• This approach involves creating a list comprehension to iterate through b and then use indexing to select values from a. It's generally less efficient for large tensors compared to torch.gather.

import torch

a = torch.arange(1, 13, dtype=torch.float).view(2, 2, 3)
b = torch.tensor([1, 0])

c = torch.tensor([a[i, idx] for i, idx in enumerate(b)])

print(c)

Custom Function with Advanced Indexing (Less Readable):

• You can create a custom function using advanced indexing techniques with boolean masks. This approach might be less readable and maintainable compared to torch.gather.

Note: Due to the complexity, we won't provide a specific code example for this method. However, it's an option if you have specific requirements or want to understand more advanced indexing concepts in PyTorch.

Choosing the Right Method:

• For most use cases, torch.gather is the recommended approach due to its efficiency and clarity.
• If you only need to select entire rows or columns based on indices with the same number of dimensions, torch.index_select can be used.
• List comprehension and custom functions might be suitable for specific scenarios or learning purposes, but they are generally less efficient or less readable.