Question

我有2个数组，一个数组的形状为(4, 2, 2)，另一个数组的形状为(4, 2).

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)

因此，对于每个0<= i <= 3，我都想拥有np.dot(A[i, :, :], B[i, :])。

如何在不使用for循环的情况下做到这一点？

Answer 1

(A @ B[...,np.newaxis])[...,0]应该起作用（np.matmul(A, B[...,np.newaxis])[...,0]的合成糖）。

对于matmul，每the documentation：

如果任一参数为N-D，N> 2，则将其视为位于最后两个索引中的一堆矩阵。

您也可以使用爱因斯坦求和np.einsum("ijk,ik->ij", A, B)，尽管在我的机器上这显然比广播matmul要慢：

import timeit

import numpy as np

A = np.random.rand(4, 2, 2)
B = np.random.rand(4, 2)


def broadcast_matmul():
    return (A @ B[..., np.newaxis])[..., 0]


def list_map():
    return np.array(list(map(np.dot, A, B)))


def einsum():
    return np.einsum("ijk,ik->ij", A, B)


print(timeit.timeit(broadcast_matmul))
print(timeit.timeit(list_map))
print(timeit.timeit(einsum))

# Results on my machine:
# 1.870921753
# 6.705698656999999
# 2.765732645

Answer 2

以下是使用np.einsum来完成此操作的方法：

np.einsum('ijk,ik->ij', A, B)

array([[0.39437083, 0.45360039],
       [0.75211742, 0.40669922],
       [0.18131254, 0.19473085],
       [0.2919673 , 0.10398859]])

使用相乘的前几个元素快速检查：

np.dot(A[0], B[0])
# array([0.39437083, 0.45360039])

Answer 3

这应该可以解决问题： list(map(np.dot, A, B))

numpy：将3D阵列与2D阵列相乘

3 个答案: