如何在PyTorch中通过不规则索引实现张量无循环赋值？

你可以利用PyTorch的广播机制生成布尔掩码，直接完成批量赋值，完全避免显式循环，效率比逐行循环高得多（尤其是当Tensor行数较多时）。

实现步骤： #

1. 将索引列表转换为PyTorch Tensor，便于后续广播操作；
2. 生成第二维度的列索引范围，通过广播比较生成掩码矩阵（标记哪些位置需要替换）；
3. 利用掩码直接对x和y的对应位置进行赋值。

代码示例： #

import torch

x = torch.randn(2, 10)
y = torch.randn(2, 10)
idxs = torch.tensor([2, 6])  # 转为Tensor适配PyTorch操作

# 生成第二维度的列索引
cols = torch.arange(x.size(1), device=x.device)
# 广播生成掩码：每行中列索引 >= 对应起始位置的元素标记为True
mask = cols >= idxs.unsqueeze(1)

# 批量赋值
x[mask] = y[mask]

另一种高级索引实现方式 #

如果你偏好直接生成需要更新的坐标，也可以通过拼接索引的方式完成：

import torch

x = torch.randn(2, 10)
y = torch.randn(2, 10)
idxs = [2, 6]

# 生成行索引：每个行号重复对应需要更新的列数
rows = torch.arange(x.size(0)).repeat_interleave([x.size(1) - idx for idx in idxs])
# 生成列索引：拼接每行需要更新的列范围
cols = torch.cat([torch.arange(idx, x.size(1)) for idx in idxs])

# 按坐标赋值
x[rows, cols] = y[rows, cols]

两种方法中，掩码实现更简洁直观，且内存效率更高，推荐优先使用。

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