如何将经Numpy swapaxes分块后的张量还原为原始输入张量？

我有如下原始4D张量（形状为 (1, 1, 6, 6)）：

([[[[ -633.,   377.,  -691.,  -683.,    50.,   936.],
    [ -121.,   851.,  -593.,  -943.,   891.,   410.],
    [ -776.,  -560.,  -270.,   816.,   425.,  -629.],
    [ -172.,   599.,  -549.,  -520.,  -684.,   -57.],
    [  265.,  -323.,   726.,   -98.,   780.,   109.],
    [   22., -1017.,    22.,   248.,  -617.,   449.]]]])

通过以下代码将其转换为池化块向量（形状为 (1, 1, 9, 4)）：

stride = 2
c = x.reshape(x.shape[0], x.shape[1], x.shape[2] // stride, stride, x.shape[3] // stride, stride).swapaxes(3, 4).reshape(x.shape[0], x.shape[1], -1, stride * stride)

转换后的池化块向量：

[[[[ -633.   377.  -121.   851.]
   [ -691.  -683.  -593.  -943.]
   [   50.   936.   891.   410.]
   [ -776.  -560.  -172.   599.]
   [ -270.   816.  -549.  -520.]
   [  425.  -629.  -684.   -57.]
   [  265.  -323.    22. -1017.]
   [  726.   -98.    22.   248.]
   [  780.   109.  -617.   449.]]]]

现在需要将这个池化块向量还原为原始输入张量，求实现方法。

还原操作是原转换步骤的逆过程，原转换依次执行了reshape→swapaxes→reshape，因此逆过程需要按相反顺序执行reshape→swapaxes→reshape，具体步骤如下：

完整还原代码（以PyTorch为例） #

import torch

# 假设c是已得到的池化块向量，x为原始张量（用于获取原始形状，若已知形状可直接指定）
stride = 2

# 第一步：将池化块向量reshape回swapaxes前的形状
temp = c.reshape(x.shape[0], x.shape[1], x.shape[2]//stride, x.shape[3]//stride, stride, stride)
# 第二步：交换轴3和轴4（逆原swapaxes操作）
temp = temp.swapaxes(3, 4)
# 第三步：reshape回原始张量形状
restored_x = temp.reshape(x.shape[0], x.shape[1], x.shape[2], x.shape[3])

# 验证还原结果与原始张量是否一致
print(torch.allclose(restored_x, x))  # 输出True，说明还原正确

步骤解释 #

1. 第一步reshape：将形状为(1,1,9,4)的池化块向量还原为(1,1,3,3,2,2)，对应原转换中最后一步reshape的逆操作；
2. 第二步swapaxes：交换第3和第4个轴，将形状变为(1,1,3,2,3,2)，抵消原转换中的swapaxes(3,4)操作；
3. 第三步reshape：将(1,1,3,2,3,2)还原为原始形状(1,1,6,6)，对应原转换中第一步reshape的逆操作。

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