如何在NumPy/PyTorch中获取3D张量2D对角线及向量化提取主对角线？

嘿，我来帮你搞定这两个需求，全用向量化实现，完全不用循环A维度，效率拉满！

问题1：在NumPy中获取3D张量的所有2D对角线 #

这里默认你的3D张量是**(A, N, N)**形状（也就是每个子矩阵都是N×N的方阵），要提取每个方阵的主对角线。

直接用NumPy内置的diagonal()方法就可以搞定，指定要取对角线的两个维度（axis1和axis2）：

import numpy as np

# 生成一个3×4×4的3D张量（3个4×4方阵）
arr = np.random.rand(3, 4, 4)
# 提取每个方阵的主对角线
diagonals = arr.diagonal(axis1=1, axis2=2)

print(diagonals.shape)  # 输出 (3, 4)，每一行对应一个方阵的主对角线

这个方法直接对整个张量做向量化操作，内部没有显式循环，速度非常快。

问题2：A×B×C张量提取每个B×C矩阵的主对角线 #

不管B和C是否相等，我们都可以用索引切片的方式实现向量化，同时支持NumPy和PyTorch。

NumPy实现 #

import numpy as np

# 生成一个3×5×3的3D张量（3个5×3的矩形矩阵）
arr = np.random.rand(3, 5, 3)
# 取B和C中较小的维度长度
min_dim = min(arr.shape[1], arr.shape[2])
# 生成对角线索引
idx = np.arange(min_dim)
# 提取每个矩阵的主对角线
diagonals = arr[:, idx, idx]

print(diagonals.shape)  # 输出 (3, 3)，每一行对应一个矩阵的主对角线

PyTorch实现 #

PyTorch的思路和NumPy完全一致，也可以用内置的diagonal()方法简化代码：

import torch

# 生成一个3×5×3的PyTorch张量
tensor = torch.rand(3, 5, 3)
# 方法1：索引切片
min_dim = min(tensor.shape[1], tensor.shape[2])
idx = torch.arange(min_dim, device=tensor.device)
diagonals = tensor[:, idx, idx]

# 方法2：用torch.diagonal更简洁
diagonals_alt = tensor.diagonal(dim1=1, dim2=2)

print(diagonals.shape)  # 输出 torch.Size([3, 3])
print(diagonals_alt.shape)  # 同样输出 torch.Size([3, 3])

为什么这种方式高效？ #

这里的索引操作是广播式的向量化操作，直接对整个张量进行切片，完全避免了手动循环A维度。当A的数值很大时，这种方法的效率会比循环高几个数量级，而且代码更简洁易读。

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