`np.dot` ohne kartesisches Produkt auf verbleibenden Achsen

Question

Gemäß den documentation:

Für N Dimensionen dot ist eine Summe Produkt in der letzten Achse der a und die zweiten bis letzten der b:

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m]) 

ich mag die Summe Produkt in der letzten Achse der a und die zweiten bis letzt der b berechnen, aber ohne das kartesische Produkt über die Form verbleibende Achsen, da die übrigen Achsen die gleiche Form haben. Lassen Sie mich mit einem Beispiel illustrieren:

a = np.random.normal(size=(11, 12, 13)) 
b = np.random.normal(size=(11, 12, 13, 13)) 
c = np.dot(a, b) 
c.shape # = (11, 12, 11, 12, 13) 

Aber ich würde die Form wie (11, 12, 13) zu sein. Der gewünschte Effekt kann mit

Ausstrahlung erreicht werden

c = np.sum(a[..., None] * b, axis=-2) 
c.shape # = (11, 12, 13) 

Aber meine Anordnungen sind relativ groß und ich möchte die Macht der parallelisierten BLAS-Implementierungen verwendet werden, die nicht angezeigt durch np.sum unterstützt werden, sondern werden von np.dot unterstützt. Irgendwelche Ideen, wie man das erreicht?

Answer 1

können Sie verwenden np.einsum -

c = np.einsum('ijk,ijkl->ijl',a,b) 

Answer 2

Sie können auch np.matmul verwenden:

c = np.matmul(a[..., None, :], b)[..., 0, :] 

Dies entspricht the new @ operator in Python 3.5 und höher:

c = (a[..., None, :] @ b)[..., 0, :] 

Es gibt nicht viel Unterschied in der Leistung - wenn überhaupt np.einsum scheint für Ihre Beispiel-Arrays marginal schneller zu sein:

In [1]: %%timeit a = np.random.randn(11, 12, 13); b = np.random.randn(11, 12, 13, 13) 
    ....: np.einsum('...i,...ij->...j', a, b) 
    ....: 
The slowest run took 5.24 times longer than the fastest. This could mean that an 
intermediate result is being cached. 
10000 loops, best of 3: 26.7 µs per loop 

In [2]: %%timeit a = np.random.randn(11, 12, 13); b = np.random.randn(11, 12, 13, 13) 
np.matmul(a[..., None, :], b)[..., 0, :] 
    ....: 
10000 loops, best of 3: 28 µs per loop