NumPy 中是否提供分层广播？

Question

我想知道是否有内置操作可以将我的代码从 Python 循环中释放出来。

问题是这样的：我有两个矩阵 A 和 B。 A 有 N 行，B 有 N 列。我想将 A 中的每个 i 行与 B 中相应的 i 列相乘（使用 NumPy 广播）。生成的矩阵将在输出中形成 i 层。所以我的结果将是 3 维数组。

NumPy 中有这样的操作吗？

Answer 1

是的，在最简单的形式中，您只需添加 "zero" 个维度，这样 NumPy 就会沿着 A 的行和 B 的列广播：

>>> import numpy as np

>>> A = np.arange(12).reshape(3, 4) # 3 row, 4 colums
>>> B = np.arange(15).reshape(5, 3) # 5 rows, 3 columns
>>> res = A[None, ...] * B[..., None]
>>> res
array([[[  0,   0,   0,   0],
        [  4,   5,   6,   7],
        [ 16,  18,  20,  22]],

       [[  0,   3,   6,   9],
        [ 16,  20,  24,  28],
        [ 40,  45,  50,  55]],

       [[  0,   6,  12,  18],
        [ 28,  35,  42,  49],
        [ 64,  72,  80,  88]],

       [[  0,   9,  18,  27],
        [ 40,  50,  60,  70],
        [ 88,  99, 110, 121]],

       [[  0,  12,  24,  36],
        [ 52,  65,  78,  91],
        [112, 126, 140, 154]]])

结果的形状为 (5, 3, 4)，如果您想要不同的形状，可以轻松移动轴。例如使用 np.moveaxis:

>>> np.moveaxis(res, (0, 1, 2), (2, 0, 1))  # 0 -> 2 ; 1 -> 0, 2 -> 1
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

形状为(3, 4, 5)。

Answer 2

直接表达您的要求的一种方法是使用 np.einsum():

>>> A = np.arange(12).reshape(3, 4)
>>> B = np.arange(15).reshape(5, 3)
>>> np.einsum('...i,j...->...ij', A, B)
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

这使用了 Einstein summation convention.

有关进一步讨论，请参阅 T. W. Körner 的 Vectors, Pure and Applied: A General Introduction to Linear Algebra 第 3 章。在里面，作者引用了爱因斯坦写给朋友的信中一段有趣的话：

"I have made a great discovery in mathematics; I have suppressed the summation sign every time that the summation must be made over an index which occurs twice..."

NumPy 中是否提供分层广播？

Is layered broadcasting available in NumPy?

python

arrays

numpy

array-broadcasting