PyTorch实现torch.gather逆操作：提取不在指定索引中的张量元素

实现逻辑 #

针对指定维度，逐行判断全量索引中没有出现在对应行idxs中的值，生成保留索引后通过gather提取结果，全程使用内置算子无自定义循环，可兼容GPU运行。

完整实现代码 #

import torch

def inverse_gather(x: torch.Tensor, idxs: torch.Tensor, dim: int = 1) -> torch.Tensor:
    """
    实现torch.gather的反向操作，提取x中指定维度未被idxs索引的元素
    要求：idxs在操作维度上每行的索引值唯一
    """
    # 生成操作维度的全量索引
    full_idx = torch.arange(x.size(dim), device=x.device, dtype=idxs.dtype).unsqueeze(0)
    # 逐行对比判断索引是否需要排除
    exclude_mask = (full_idx.unsqueeze(-1) == idxs.unsqueeze(1)).any(dim=-1)
    # 生成保留的索引并提取结果
    keep_idx = full_idx.expand_as(exclude_mask)[~exclude_mask].view(x.size(0), x.size(dim) - idxs.size(dim))
    return torch.gather(x, dim, keep_idx)

测试验证 #

# 输入示例
x = torch.arange(30).reshape(3,10)
idxs = torch.tensor([[1,2,3], [4,5,6], [7,8,9]], dtype=torch.long)
# 调用函数
output = inverse_gather(x, idxs)
print(output)

输出结果：

tensor([[ 0,  4,  5,  6,  7,  8,  9],
        [10, 11, 12, 13, 17, 18, 19],
        [20, 21, 22, 23, 24, 25, 26]])

方案优势 #

• 无任何自定义for循环，所有运算均为PyTorch内置向量化算子，运行效率高，可直接在GPU上执行
• 支持idxs为torch.topk输出的场景，满足每行索引唯一的前置要求
• 调整dim参数即可适配不同维度的反向gather需求

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