оптимизируйте создание ndarray, содержащих матрицы, с использованием тензордота или einsum

#python-3.x #numpy

Вопрос:

Я хочу создать массив матриц , умножив матрицу (n, n) на элементы из другого массива формы (i, j) , ожидаемая результирующая форма должна быть (i, j, n, n) .

 [[a1, a2], [b1, b2], ...] X M = [[a1 * M, a2 * M], [b1 * M, b2 * M], ...]  

Оказывается, это своего рода медленная операция, и я ищу способы ее оптимизации. В настоящее время я использую

 np.tensordot(a, b, axes=0)  

что немного быстрее, чем версия einsum

 np.einsum('ij,ab-gt;ijab', a, b, optimize=True)  

Эти выражения оказываются узкими местами, можно ли выразить их по-другому, чтобы повысить скорость? Для фигур (50, 100531) (11, 11) в настоящее время это занимает ~3 секунды.

Комментарии:

1. a[:,:,None,None]*b