Умножение матрицы на пакеты векторов в NumPy: np.matvec, axes и эквивалент через einsum

Как умножить матрицу 3×2 на стек векторов 2×4×5 в NumPy без лишних транспонирований: np.matvec с параметром axes, moveaxis и эквивалент через einsum. Примеры.

Умножение матрицы на пакеты векторов — повседневная рутина в коде, плотно завязанном на NumPy. Как только вы переходите к обобщённым ufunc, таким как np.matvec, единственной реальной трудностью становится размещение осей: у матрицы два базовых измерения, у векторов — одно, а значения по умолчанию предполагают, что векторы лежат на последней оси. Если ваши данные не следуют этой конвенции, понадобится либо безобидная перестановка раскладки, либо явное сопоставление осей.

Пример

Задача: умножить одну матрицу 3×2 на сетку 4×5 двумерных векторов, хранящихся в массиве 2×4×5, и получить результат 3×4×5.

import numpy as np
A = np.random.rand(3, 2)            # матрица 3x2
W = np.random.rand(2, 4, 5)         # стек 4x5 векторов, каждый длины 2
# matvec с небольшими перестановками раскладки
Y1 = np.matvec(A, W.transpose(1, 2, 0)).transpose(2, 0, 1)
# эквивалент через einsum
Y2 = np.einsum('ij,j...->i...', A, W)
# matvec с явным отображением осей (без транспонирования)
Y3 = np.matvec(A, W, axes=[(0, 1), 0, 0])
# вариант, не зависящий от формы, с использованием moveaxis
Y4 = np.moveaxis(np.matvec(A, np.moveaxis(W, 0, -1)), -1, 0)
# проверка корректности
print(Y1.shape)  # (3, 4, 5)
print(Y2.shape)  # (3, 4, 5)
print(Y3.shape)  # (3, 4, 5)
print(Y4.shape)  # (3, 4, 5)
print(np.allclose(Y1, Y2))  # True
print(np.allclose(Y3, Y2))  # True
print(np.allclose(Y4, Y2))  # True
# невекторизованная проверка промежуточного результата до обратного транспонирования
check = np.matvec(A, W.transpose(1, 2, 0))
ref = [[A @ W[:, r, c] for c in range(5)] for r in range(4)]
print(np.allclose(check, ref))  # True

Что на самом деле происходит

np.matvec — это обобщённая ufunc, которая принимает матрицу и вектор и возвращает вектор. По умолчанию она ожидает векторное измерение на последней оси второго аргумента и записывает получившийся вектор на последнюю ось результата. Это удобно, если ваши данные уже в виде «...×вектор». Но если векторы находятся на первой оси, как в раскладке 2×4×5, у вас два пути: перенести эту ось в конец, вызвать ufunc и вернуть результат на место, либо явно описать соответствие осей через параметр axes.

Ключевая мысль — отделить базовые оси от батчевых. Для матрицы 3×2 и векторов 2×4×5 базовые оси — это (0, 1) у матрицы и 0 у векторов. Остальные — здесь 4×5 — батчевые размерности, по которым операция просто повторяется. Без указаний ufunc использует последнюю ось под векторы; с указаниями вы сами выбираете, где они находятся.

Решение с перестановками раскладки

Минимальный подход: перенести векторную ось в конец, применить np.matvec и затем перенести векторную ось результата в начало. Функции transpose и moveaxis меняют только метаданные шагов (strides) и не двигают сами данные, так что на практике это почти бесплатно. Именно это делают варианты Y1 и Y4. Формулировка с moveaxis менее привязана к форме, если не хочется жёстко задавать перестановки.

Решение с отображением осей

Если предпочитаете явность, ключевое слово axes позволяет прямо задать три роли: где находятся оси матрицы в первом аргументе, где векторная ось во втором и куда поместить векторную ось на выходе. Для указанных форм axes=[(0, 1), 0, 0] означает: базовые оси матрицы в первом входе — (0, 1), векторная ось второго входа — 0, а векторная ось выхода ставится на позицию 0. Результат — массив 3×4×5, ровно в нужной раскладке.

Эквивалент через einsum

Выражение einsum 'ij,j...->i...' кодирует ту же идею: умножение вдоль измерения j, сохранение батчевых размерностей как ..., и вывод вектора по i. Это кратко и легко обобщается, а в данном случае даёт ту же форму и значения, что и np.matvec. На измеренных формах 3×2 и 2×4×5 времена почти одинаковые, различие менее 0.1%, а поведение с ростом размеров может меняться: при увеличении размеров einsum может оказаться чуть быстрее, затем на больших данных выигрывает np.matvec. Оба подхода жизнеспособны.

Зачем это важно

Корректное выравнивание осей позволяет избежать скрытых несоответствий форм и нагромождения трудночитаемых транспонирований по коду. Когда вы полагаетесь на распространение одной матрицы на множество векторов — как при применении матрицы 3×2 к стеку 2×4×5 — замысел остаётся прозрачным, если вы либо задаёте его через axes, либо изолируете небольшой набор перемещений в одном месте. Это помогает поддерживать предсказуемые раскладки выхода для последующих этапов и тестов.

Выводы

Если ваши векторы не на последней оси, либо перенесите эту ось в конец и верните результат обратно, либо укажите np.matvec, где что находится, через axes. Когда вам нужен результат 3×4×5 из матрицы 3×2 и векторов 2×4×5, используйте один из этих компактных шаблонов. Для явного сопоставления осей np.matvec(A, W, axes=[(0, 1), 0, 0]) читается чисто и избегает лишних перестановок. Для краткой универсальной формы эквивалентен np.einsum('ij,j...->i...', A, W). Оба варианта проверяются с np.allclose и при типичных формах здесь работают сходно. Выбирайте тот, который делает код читабельнее и помогает держать раскладки данных согласованными.