在两个数组的最终轴上逐元素相乘

Question

给定一个 3d 数组和一个 2d 数组，

a = np.arange(10*4*3).reshape((10,4,3))
b = np.arange(30).reshape((10,3))

我如何运行跨每个最终轴的元素乘法，导致 c 其中 c 的形状 .shape 为 a ？ IE。

c[0] = a[0] * b[0]
c[1] = a[1] * b[1]
# ...
c[i] = a[i] * b[i]

Answer 1

使用np.einsum:

c = np.einsum('ijk,ik->ijk', a, b)

快速检查：

print(np.allclose(c[0], a[0] * b[0]))
print(np.allclose(c[1], a[1] * b[1]))
print(np.allclose(c[-1], a[-1] * b[-1]))
# True
# True
# True

Answer 2

在不涉及任何总和减少的情况下，使用 np.newaxis/None -

将 b 扩展到 3D 后，简单的 broadcasting 将非常有效

a*b[:,None,:] # or simply a*b[:,None]

运行时测试 -

In [531]: a = np.arange(10*4*3).reshape((10,4,3))
     ...: b = np.arange(30).reshape((10,3))
     ...: 

In [532]: %timeit np.einsum('ijk,ik->ijk', a, b) #@Brad Solomon's soln
     ...: %timeit a*b[:,None]
     ...: 
100000 loops, best of 3: 1.79 µs per loop
1000000 loops, best of 3: 1.66 µs per loop

In [525]: a = np.random.rand(100,100,100)

In [526]: b = np.random.rand(100,100)

In [527]: %timeit np.einsum('ijk,ik->ijk', a, b)
     ...: %timeit a*b[:,None]
     ...: 
1000 loops, best of 3: 1.53 ms per loop
1000 loops, best of 3: 1.08 ms per loop

In [528]: a = np.random.rand(400,400,400)

In [529]: b = np.random.rand(400,400)

In [530]: %timeit np.einsum('ijk,ik->ijk', a, b)
     ...: %timeit a*b[:,None]
     ...: 
10 loops, best of 3: 128 ms per loop
10 loops, best of 3: 94.8 ms per loop

在两个数组的最终轴上逐元素相乘

Multiplying elementwise over final axis of two arrays

python

numpy

python-3.x

numpy-einsum