Объединение элемента -мудрого и матричного умножения с многомерными -массивами в NumPy

У меня есть два многомерных массива NumPy, Aи B, с A.shape = (K, d, N)и B.shape = (K, N, d). Я хотел бы выполнить мудрую операцию элемента -над осью 0(K), при этом эта операция представляет собой умножение матриц по осям 1 и 2(d, NиN, d). Таким образом, результатом должен быть многомерный массив Cс C.shape = (K, d, d), так что C[k] = np.dot(A[k], B[k]). Простая реализация выглядела бы так :

C = np.vstack([np.dot(A[k], B[k])[np.newaxis, :, :] for k in xrange(K)])

, но эта реализация медленная . Немного более быстрый подход выглядит следующим образом :

C = np.dot(A, B)[:, :, 0, :]

, который использует поведение np.dotпо умолчанию для многомерных массивов, что дает мне массив с формой (K, d, K, d). Однако этот подход вычисляет требуемый ответ K, умноженный на (, каждый из элементов вдоль оси 2 одинаков ). Асимптотически это будет медленнее, чем первый подход, но накладные расходы намного меньше. Я также знаю о следующем подходе :

from numpy.core.umath_tests import matrix_multiply
C = matrix_multiply(A, B)

, но я не гарантирую, что эта функция будет доступна. Таким образом, мой вопрос заключается в том, обеспечивает ли NumPy стандартный способ сделать это эффективно? Ответ, применимый к многомерным массивам в целом, был бы идеальным, но ответ, специфичный только для этого случая, тоже был бы замечательным.

Редактировать:Как указал @Juh _, второй подход неверен. Правильная версия — :

C = np.dot(A, B).diagonal(axis1=0, axis2=2).transpose(2, 0, 1)

, но добавленные накладные расходы делают ее медленнее , чем первый подход, даже для небольших матриц. Последний подход выигрывает с большим отрывом во всех моих временных тестах для малых и больших матриц. Теперь я серьезно рассматриваю возможность использования этого, если не появится лучшего решения, даже если это будет означать копирование библиотеки numpy.core.umath_tests(, написанной на C ), в мой проект.