如何用一维张量替代标量对二维张量进行切片？

在PyTorch里，常规的t[:, :k]只能给所有行统一设置切片长度，但要实现每行按自定义长度切片，用掩码处理是最直接的方案，具体实现如下：

实现思路 #

• 生成对应每个元素的行、列索引，通过广播机制让索引矩阵和原张量形状匹配
• 创建布尔掩码，标记每行中需要保留的元素（列索引小于该行指定的切片长度）
• 初始化结果张量，用NaN或0填充未保留的位置，再将掩码标记的区域替换为原张量的值

代码示例 #

import torch

# 初始化输入张量（修正原示例的张量创建语法错误）
t = torch.tensor([[1, 1, 1], [2, 2, 2], [3, 3, 3]], dtype=torch.float32)
# 每行的目标切片长度，转为整数张量
k = torch.tensor([1, 2, 3])

# 生成行、列索引矩阵
rows = torch.arange(t.shape[0]).unsqueeze(1)  # 形状 (3,1)
cols = torch.arange(t.shape[1]).unsqueeze(0)  # 形状 (1,3)

# 生成掩码：判断每个位置的列索引是否小于对应行的切片长度
mask = cols < k[rows]

# 初始化结果张量，用NaN填充（替换为0即可实现0掩码）
result = torch.full_like(t, float('nan'))
# 保留掩码标记的元素
result[mask] = t[mask]

print(result)

输出结果 #

tensor([[1., nan, nan],
        [2., 2., nan],
        [3., 3., 3.]])

内容的提问来源于stack exchange，提问作者David Davíd