转置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)