在numpy ndarray中重新排列和重塑元素

我有一个由外部模块返回的numpy数据数组。阵列的形状为（3,3128）。基本上是一堆128个瓷砖，每个瓷砖为3x3

---

如何重新排序，使形状成为（128,3,3）。通过这种方式，可以更容易地按磁贴编号进行索引。最后一步是展平到（128,9），这样128个分片中的每个分片都可以作为9值向量轻松访问。

您可以使用指定新数组顺序的转置，例如

a = np.arange(0,3*3*128).reshape(3,3,128) 
a_reorder = a.transpose([2,0,1])

您可以通过比较所有瓷砖来检查其是否正确

np.all([np.all(a[:,:,i]==a_reorder[i,:,:]) for i in range(128)])

并用

a_flat = a_reorder.reshape(128,9)

---

将3*3*128重塑为128*3*3：

y = einops.rearrange(x, 'x y tile -> tile x y')

y = einops.rearrange(x, 'x y tile -> tile (x y)')

或者我们可以在一次操作中直接重塑为128*9