使用_scatter()替换矩阵元素：基于子序列长度排序张量x的问题

我来帮你搞定这个问题！你之前用scatter_没得到正确结果，是因为这个函数的作用是写入值，而你需要的是根据排序索引提取并重新排列元素，这时候应该用torch.gather才对。

问题分析与解决方案 #

你的需求很明确：每个batch里，把x的sequence_length维度的元素，按照y对应位置的子序列长度降序排列。我们一步步来实现：

1. 对y降序排序，获取排序索引
   这一步你已经做对了，y_sort_idx里存的就是每个batch中，原y元素按降序排列对应的原索引位置：

   y_sorted, y_sort_idx = y.sort(dim=1, descending=True)
   

   比如第二个batch的y是[2,1,2]，降序排序后对应的原索引是[0,2,1]——这正是我们用来重新排列x的关键依据。

2. 用gather重新排列x
   因为x的形状是[batch_size, sequence_length, subseq_length]，而y_sort_idx是[batch_size, sequence_length]，我们需要把索引扩展成和x相同的维度（在最后加一个维度并广播），这样才能对应到每个子序列的元素。

完整代码示例 #

import torch

x = torch.tensor([[[1, 2], [2, 0], [0, 0]], [[2, 2], [2, 0], [3, 3]]])
y = torch.tensor([[2, 1, 0], [2, 1, 2]])

# 对y降序排序，得到排序后的索引
y_sorted, y_sort_idx = y.sort(dim=1, descending=True)

# 扩展索引维度使其与x匹配，通过gather提取并排列元素
x_sorted = x.gather(dim=1, index=y_sort_idx.unsqueeze(-1).expand_as(x))

print(x_sorted)

运行后输出就是你期望的结果：

tensor([[[1, 2], [2, 0], [0, 0]],
        [[2, 2], [3, 3], [2, 0]]])

为什么scatter_不行？ #

简单来说：

• scatter_的逻辑是把src的值写入到目标张量的指定位置，你用x本身作为目标时，相当于在原张量上覆盖写入，不符合“按索引提取元素重新排列”的需求；
• gather的逻辑是从原张量的指定位置提取元素，正好匹配我们要根据排序索引重新组织x的场景。

补充说明 #

y_sort_idx.unsqueeze(-1).expand_as(x)的作用：

• unsqueeze(-1)把y_sort_idx从[2,3]变成[2,3,1]；
• expand_as(x)把它扩展成和x一样的[2,3,2]，让每个索引都能对应到子序列的每一个元素，保证整个子序列被正确提取。

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