PyTorch中高效提取非对称张量的特定对角线元素

你有一个形状为(5 * n, n)的PyTorch张量，需要按以下规则提取元素：第一列取前5行，第二列取接下来的5行，以此类推处理完所有列。下面是两种无循环的高效实现方式，适合大规模张量的重复运算：

示例输入 #

import torch
t = torch.arange(45).reshape(15, 3)

期望输出 #

tensor([0, 3, 6, 9, 12, 16, 19, 22, 25, 28, 32, 35, 38, 41, 44])

解决方案 #

方法1：重塑+高级索引 #

把原张量拆分成n个5行n列的子块，每个子块取对应位置的列，最后展平结果：

n = t.shape[1]
out = t.reshape(n, 5, n)[torch.arange(n), :, torch.arange(n)].flatten()

方法2：构造索引对直接取值 #

生成匹配的行、列索引，直接通过张量索引提取目标元素：

n = t.shape[1]
row_indices = torch.arange(5 * n)
col_indices = torch.arange(n).repeat_interleave(5)
out = t[row_indices, col_indices]

两种方法都是纯张量操作，没有Python循环，运算效率极高，适合上千次重复处理大张量的场景。

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