转置3D数组和多元素内存邻接性

Question

我有一个巨大的3D数组，看起来像A.shape = (100000, 5000, 50)。我需要将其转置为A.shape = (50, 5000, 100000)形式的数组。然后，我需要对A中包含的50个矩阵中的每个矩阵执行a = a.T @ a操作。这给出了一个形式为A.shape = (50, 5000, 5000)的3D数组。

如果我用A.transpose(2, 1, 0) @ A.transpose(2, 0, 1)这样做，单矩阵乘法a = a.T @ a会比a不是从A中提取出来的情况慢一千倍。

问题是，在转置之后，3D数组不是连续的。我试着在转置后使用np.ascontiguousarray()或copy()。它得到了改进，但仍然较慢，并且需要花费相当多的时间进行复制。

有没有人能建议一个更好的选择？特别是，我正在尝试使用np.einsum，但是我不能。

Answer 1

您可以尝试以下操作：

A = ...
b = np.einsum('jki,jli->ikl', A, A)
print(b.shape)
# (50, 5000, 5000)