Точечный продукт вещает на 3D сетке с NumPy

Question

Я пытаюсь понять, как транслировать систематические операции с точечным продуктом на трехмерной сетке 10x10x10. Я сформулировал следующие массивы:

A.shape=(5,5,10,10,10)
b.shape=(5,10,10,10)

Я хочу получить массив, подобный следующему

c.shape=(5,10,10,10)

Который я до сих пор получил через следующий код

c=np.sum(A*b,axis=1)

Я считаю, однако, что я смогу получить тот же результат с помощью np.dot или np.tensordot. Я старался изо всех сил, но мне не удалось получить эквивалентный результат. Это очень помогло бы мне понять, как работает np.tensordot, так как в дальнейшем мне также потребуется использовать np.tensorsolve.

Answer 1

Если «точечные» размеры находятся в конце, matmul будет работать

Сравнение 3 методов:

In [252]: A=np.arange(5*5*10*10*10).reshape((5,5,10,10,10))
In [253]: B=np.arange(5*10*10*10).reshape((5,10,10,10))

In [254]: C=np.sum(A*B, axis=1)
In [255]: D=np.einsum('ijklm,jklm->iklm',A,B)
In [256]: E = (A.transpose(2,3,4,0,1)@B.transpose(1,2,3,0)[...,None])[...,0].transpose(3,0,1,2)

Все транспонирования превращают массивы в (....,5,5) и (...,5,1).

In [257]: np.allclose(C,D)
Out[257]: True
In [258]: np.allclose(C,E)
Out[258]: True

In [259]: timeit C=np.sum(A*B, axis=1)
124 µs ± 4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [260]: timeit D=np.einsum('ijklm,jklm->iklm',A,B)
66 µs ± 18.7 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [261]: timeit E = (A.transpose(2,3,4,0,1)@B.transpose(1,2,3,0)[...,None])[...
     ...: ,0].transpose(3,0,1,2)
68.6 µs ± 973 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

tensordot изменяет форму и транспонирует массивы, чтобы он мог выполнить простое 2d dot - и затем конвертирует обратно. Последние 3 измерения фактически являются одним измерением 1000.

In [262]: np.tensordot(A,B,(1,0)).shape
Out[262]: (5, 10, 10, 10, 10, 10, 10)
In [263]: timeit np.tensordot(A,B,(1,0)).shape
74 ms ± 70.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Результат намного больше - своего рода внешнее произведение в не суммирующих измерениях. Результат там, похоронен как многомерная диагональ.

Тензорот эффективно:

In [264]: X=(B.reshape(5,1000).T).dot(A.reshape(5,5,1000))
In [265]: X.shape
Out[265]: (1000, 5, 1000)
In [266]: timeit X=(B.reshape(5,1000).T).dot(A.reshape(5,5,1000))
78.9 ms ± 82.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

Нам нужно сохранить последние четыре оси выровненными и иметь их на выходе, кроме второй оси (ось = 1), которая должна быть уменьшена до суммы. Для такого случая np.einsum - это путь, кроме np.matmul. С np.einsum это было бы проще / интуитивно понятно, как это -

c = np.einsum('ijklm,jklm->iklm',A,b)