numpy:在两个2d数组的一个公共轴上进行广播乘法

问题描述

我正在寻找一种以元素方式分别乘以形状(a，b)和(b，c)的两个2d数组的方法.在"b"轴上，这是两个数组的共同点.

I'm looking for a way to element-wise multiply two 2d arrays of shape (a, b) and (b, c), respectively. Over the 'b' axis, which the two arrays have in common.

例如，我要广播(向量化)的示例是:

For instance, an example of what I'd like to broadcast (vectorize) is:

import numpy as np    

# some dummy data
A = np.empty((2, 3))
B = np.empty((3, 4))

# naive implementation
C = np.vstack(np.kron(A[:, i], B[i, :]) for i in [0, 1, 2])

# this should give (3, 2, 4)
C.shape

有人知道在这里做什么吗?有更好的方法吗?

Does anyone know what to do here? Is there a better way?

推荐答案

A和B
的定义，请注明@hpaulj使用np.outer和np.stack

credit to @hpaulj for the definitions of A and B
use np.outer and np.stack

A = np.arange(6).reshape((2, 3))
B = np.arange(12).reshape((3, 4))

np.stack([np.outer(A[:, i], B[i, :]) for i in range(A.shape[1])])

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

并获得正确形状的np.einsum

np.einsum('ij, jk->jik', A, B)

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

广播和transpose

(A[:, None] * B.T).transpose(2, 0, 1)

[[[ 0  0  0  0]
  [ 0  3  6  9]]

 [[ 4  5  6  7]
  [16 20 24 28]]

 [[16 18 20 22]
  [40 45 50 55]]]

---

形状为(3, 2, 4)

定时

timing

这篇关于numpy:在两个2d数组的一个公共轴上进行广播乘法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！