如何矢量化将函数传递给两个 numpy 数组：3D 和 2D？

Question

我有两个多维 numpy 数组：x 是 3D，y 是 2D。

如果我有一个函数foo(a, b)，它将两个二维数组作为输入，我如何将我的多维数组传递给foo，并以矢量化方式迭代x的第三维以获得一个foo的结果列表？

我一直在尝试使用 np.vectorize 执行此操作，但它会遍历数组的行并产生错误，所以我被卡住了。

Answer 1

您可以使用signature 关键字指定函数的签名。但是，这将尝试使用每个输入的最后一个维度，因此您必须手动转置。示例

F = np.vectorize(np.matmul, signature='(m,n),(n,l)->(m,l)', otypes=(float,))
A = np.arange(12).reshape(2, 2, 3)
B = np.diag((1.5, 2.5))
# F(A.transpose(2,0,1), B)
# array([[[ 0. ,  7.5],
#         [ 9. , 22.5]],
# 
#        [[ 1.5, 10. ],
#         [10.5, 25. ]],
#
#        [[ 3. , 12.5],
#         [12. , 27.5]]])

正如@hpaulj 在 cmets 中指出的那样，vectorize 是一种便利功能，而不是性能增强器。