通过使用 np.einsum 将 2D 数组的行广播到 3D 数组的每个切片来进行乘法

Question

我有数组A 和B：

>>> import numpy as np

>>> A = np.ones((3,3,2))

>>> B = np.array([
    [0,0],
    [1,1],
    [2,2],
])

我想将B的每一行乘以A的每一切片，这样B的每一行就在A的每一切片上进行广播，即：

>>> np.array([A_slice*B_row for A_slice, B_row in zip(A, B)])
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

我想要最有效的方法，我相信这可能是使用np.einsum（但是，如果您认为使用另一种方法更快，例如我在下面提到的方法，请告诉我）。

我尝试了以下方法：

>>> np.einsum('ijk,lk->ijk', A, B)
[[[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]]

如您所见，这显然与上面的输出不同。

我能想到的另一个解决方案是：

>>> A*B[:,np.newaxis,:].repeat(3, axis=1)
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

哪个确实给出了正确的输出，但我仍然很想知道如何使用np.einsum 来做到这一点

编辑： Warren Weckesser 在 cmets 中指出，上面的解决方案可以简化为 A*B[:,np.newaxis,:]，这是我见过的不使用 np.einsum 的最干净的解决方案。

Answer 1

numpy.einsum 解决方案：

C = np.einsum('ijk,jk->jik', A, B)

使用省略号：

C = np.einsum('ij...,j...->ji...', A, B)

输出

[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]