PyTorch批量块对角矩阵构造实现方案咨询

针对批量张量构造块对角矩阵的需求，这里有两种简单高效的实现方式：

方法一：直接赋值填充零矩阵 #

先初始化对应形状的全零张量，再将两个输入张量分别填充到结果的左上角和右下角区域：

import torch

# 示例参数
n, d, e = 3, 2, 3
A = torch.randn(n, d, d)
B = torch.randn(n, e, e)

# 创建结果张量
result = torch.zeros(n, d+e, d+e, dtype=A.dtype, device=A.device)
# 填充A到左上角区块
result[:, :d, :d] = A
# 填充B到右下角区块
result[:, d:, d:] = B

方法二：Pad后相加 #

通过对两个张量分别做padding扩展到目标形状，再相加得到块对角矩阵：

import torch
import torch.nn.functional as F

n, d, e = 3, 2, 3
A = torch.randn(n, d, d)
B = torch.randn(n, e, e)

# 对A做右、下方向的padding，补e个零
A_padded = F.pad(A, (0, e, 0, e))
# 对B做左、上方向的padding，补d个零
B_padded = F.pad(B, (d, 0, d, 0))
# 相加得到块对角结果
result = A_padded + B_padded

这两种方法都能高效处理批量场景，避开了torch.block_diag仅支持二维张量的局限。

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