PyTorch技术问答：如何用给定排序索引重排张量

问题说明 #

现有两个张量：

• 张量A，形状为(d₀, d₁, ..., dₙ, dₙ₊₁)
• 排序索引张量I，形状为(d₀, d₁, ..., dₙ)

需要利用I中的索引对A的指定维度进行重排：A与I的前n个维度完全一致，A的第n+1维度长度无限制，重排操作针对A的第n个维度（即I最后一个维度对应的A的维度）。

示例 #

当A的形状为torch.Size([8,8,4])，I的形状为torch.Size([8,8])时，需将A的倒数第二维度（从0开始计数为第1维）的元素，按照I的索引顺序重新排列。

解决方案 #

以下是支持任意维度张量的通用重排函数，通过调整索引张量的形状以匹配目标张量，再借助torch.gather完成重排操作：

import torch
from torch import LongTensor, Tensor

def sort_by_indices(values: Tensor, indices: LongTensor) -> Tensor:
    num_dims = indices.dim()

    # 调整索引张量形状，新增维度以匹配values的维度数量
    new_shape = tuple(indices.shape) + tuple(1 for _ in range(values.dim() - num_dims))
    # 定义各维度的重复次数，让索引张量在values的额外维度上完成复制
    repeats = tuple(1 for _ in range(num_dims)) + tuple(values.shape[num_dims:])

    # 扩展索引张量的形状，使其与values完全匹配
    repeated_indices = indices.reshape(*new_shape).repeat(*repeats)
    # 在指定维度执行gather操作，完成重排
    return torch.gather(values, num_dims - 1, repeated_indices)

内容的提问来源于stack exchange，提问作者mattiavandi