Как векторизовать передачу функции в два массива: 3D и 2D?

Question

У меня есть два многомерных массива numpy: x - это 3D, а y - это 2D.

Если у меня есть функция foo(a, b), которая принимает в качестве входных данных два двумерных массива, как я могу это сделать?перейти к foo моим многомерным массивам и перебрать 3-е измерение 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 - это удобная функция, а не повышение производительности.