如何使用 numpy（而不是 for 循环）将矩阵的矩阵元素相乘？

Question

我有一些 4 维 numpy 数组，其中最简单的可视化是任意大小的矩阵（不一定是正方形），其中每个元素都是 2x2 方阵。 我想将标准矩阵乘以（@）大矩阵的 2x2 矩阵元素（产生另一个相同维度的 2x2 矩阵矩阵）。 最终的希望是使用 CuPy 并行化这个过程，所以我想这样做，而不是求助于循环更大矩阵的每个元素。

任何帮助，将不胜感激。

我的意思的例子：

x = np.array([[  [[1,0], [0, 1]], [[2,2], [2, 1]]  ]])
y = np.array([[  [[1,3], [0, 1]], [[2,0], [0, 2]]  ]])
xy = np.array([[  [[1,3], [0, 1]], [[4,4], [4, 2]]  ]])


[[ [[1, 0],            [[2, 2]          x         [[ [[1, 3],            [[2, 0]
    [0, 1]]      ,      [2, 1]] ]]                    [0, 1]]      ,      [0, 2]] ]]




=> [[ [[1, 3],            [[4, 4]
       [0, 1]]      ,      [4, 2]] ]]

在此示例中，2 个“大”矩阵是 1x2 矩阵，其中 2 个元素中的每一个都是 2x2 矩阵。 我试图以一种清楚地说明正在发生的事情以及使用标准 4d numpy 数组的方式对其进行布局。

根据评论编辑。

Answer 1

正如 Homer512 在评论中所说， np.matmul （又名@运算符）将处理这种情况（请参阅numpy 文档）。 您需要确保您的 2 x 2 矩阵在最后一个维度中。

import numpy as np

a1 = np.array([[1, 0], [0, 1]])
a2 = np.array([[2, 2], [2, 1]])
a = np.array([a1, a2])

b1 = [[1, 3], [0, 1]]
b2 = [[2, 0], [0, 2]]
b = np.array([b1, b2])

x = np.array([a, b])

print(a @ b)

输出：

[[[1 3]
  [0 1]]

 [[4 4]
  [4 2]]]