PyTorch中按指定索引将张量A的值投影到张量B的实现需求

问题描述 #

给定以下PyTorch张量：

import torch

A = torch.tensor([[ 0.4821, -0.3484,  0.0915, -0.1870],
                  [ 1.3817,  0.3011,  1.0704,  2.1717]])

B = torch.zeros(2,6)

C = torch.tensor([[1,2,2,3], [3,7,2,5]])  # 形状与A相同

需要完成以下操作：

• 仅处理C中小于6的索引值（B的最后一维长度为6，索引范围0-5）
• 将B中对应这些索引的位置，替换为A的对应值
• 若同一索引对应多个A的值，取最大值（示例输出采用该逻辑）

实现代码 #

import torch

# 初始化张量
A = torch.tensor([[ 0.4821, -0.3484,  0.0915, -0.1870],
                  [ 1.3817,  0.3011,  1.0704,  2.1717]])
B = torch.zeros(2,6)
C = torch.tensor([[1,2,2,3], [3,7,2,5]])

# 1. 筛选出C中符合要求的索引，生成有效掩码
valid_mask = C < 6
# 提取有效索引、对应A的值，以及行索引
valid_indices = C[valid_mask]
valid_A_values = A[valid_mask]
row_indices = torch.arange(A.size(0)).repeat_interleave(A.size(1))[valid_mask]

# 2. 用scatter_reduce实现同一索引下取最大值的赋值逻辑
B.scatter_reduce_(dim=1, index=valid_indices.unsqueeze(1), src=valid_A_values.unsqueeze(1), reduce='max')

print("处理后的B张量：")
print(B)

代码解释 #

• 有效掩码过滤：通过C < 6排除掉C中像7这类超出B索引范围的值，只保留有效索引
• 元素提取：从A和C中提取有效位置的数值与索引，同时生成行索引确保每个值对应到B的正确行
• 聚合赋值：使用scatter_reduce_方法，指定reduce='max'实现同一索引下取最大值的逻辑，直接原地修改B张量

运行代码后输出与期望一致：

处理后的B张量：
tensor([[0.0000, 0.4821, 0.0915, -0.1870, 0.0000, 0.0000],
        [0.0000, 0.0000, 1.0704, 1.3817, 0.0000, 2.1717]])

如果需要改为求和逻辑，仅需将reduce='max'替换为reduce='sum'即可。

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