numpy: multiplicación de difusión sobre un eje común de dos matrices 2d

Estoy buscando una manera de multiplicar por dos elementos dos matrices de formas 2d (a, b) y (b, c), respectivamente. Sobre el eje ‘b’, que las dos matrices tienen en común.

Por ejemplo, un ejemplo de lo que me gustaría transmitir (vectorizar) es:

import numpy as np # some dummy data A = np.empty((2, 3)) B = np.empty((3, 4)) # naive implementation C = np.vstack(np.kron(A[:, i], B[i, :]) for i in [0, 1, 2]) # this should give (3, 2, 4) C.shape 

¿Alguien sabe qué hacer aquí? ¿Hay alguna manera mejor?

Crédito a @hpaulj por las definiciones de A y B
usa np.outer y np.stack

 A = np.arange(6).reshape((2, 3)) B = np.arange(12).reshape((3, 4)) np.stack([np.outer(A[:, i], B[i, :]) for i in range(A.shape[1])]) [[[ 0 0 0 0] [ 0 3 6 9]] [[ 4 5 6 7] [16 20 24 28]] [[16 18 20 22] [40 45 50 55]]] 

y para obtener el np.einsum en la forma correcta

 np.einsum('ij, jk->jik', A, B) [[[ 0 0 0 0] [ 0 3 6 9]] [[ 4 5 6 7] [16 20 24 28]] [[16 18 20 22] [40 45 50 55]]] 

Difusión y transpose

 (A[:, None] * BT).transpose(2, 0, 1) [[[ 0 0 0 0] [ 0 3 6 9]] [[ 4 5 6 7] [16 20 24 28]] [[16 18 20 22] [40 45 50 55]]] 

la forma es (3, 2, 4)

sincronización
introduzca la descripción de la imagen aquí

Con diferentes casos de prueba:

 In [56]: A=np.arange(6).reshape((2,3)) In [57]: B=np.arange(12).reshape((3,4)) In [58]: np.vstack([np.kron(A[:,i],B[i,:]) for i in range(3)]) Out[58]: array([[ 0, 0, 0, 0, 0, 3, 6, 9], [ 4, 5, 6, 7, 16, 20, 24, 28], [16, 18, 20, 22, 40, 45, 50, 55]]) 

Un primer bash con `einsum, conservando los 3 ejes (sin sumr)

 In [60]: np.einsum('ij,jk->ijk',A,B) Out[60]: array([[[ 0, 0, 0, 0], [ 4, 5, 6, 7], [16, 18, 20, 22]], [[ 0, 3, 6, 9], [16, 20, 24, 28], [40, 45, 50, 55]]]) 

Los mismos números, pero en una forma diferente.

Puedo reordenar los ejes en la salida, para hacer un 2x4x3, que se puede remodelar a 8,3 y transponer.

 In [64]: np.einsum('ij,jk->ikj',A,B).reshape(8,3).T Out[64]: array([[ 0, 0, 0, 0, 0, 3, 6, 9], [ 4, 5, 6, 7, 16, 20, 24, 28], [16, 18, 20, 22, 40, 45, 50, 55]]) 

Así que con otra iteración puedo deshacerme de la transposición.

 In [68]: np.einsum('ij,jk->jik',A,B).reshape(3,8) Out[68]: array([[ 0, 0, 0, 0, 0, 3, 6, 9], [ 4, 5, 6, 7, 16, 20, 24, 28], [16, 18, 20, 22, 40, 45, 50, 55]]) 

Debería haber llegado allí de inmediato. A es (2,3), B es (3,4) y quiero (3,2,4) remodelado a (3,8). i = 2, j = 3, k = 4 => jik.

Así que otra forma de describir el problema,

 a_ij * b_jk = c_jik 

Y como no uso la sum del einsum , la multiplicación regular también funcionará, con uno o más transpuestos.