如何将函数传递给两个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在注释中指出的，vectorize是一个方便的函数，而不是性能增强器。你知道吗