Numpy elementowo produkt z tablicy 3d

Question

Mam dwie tablice 3d A i B o kształcie (N, 2, 2), które chciałbym pomnożyć zgodnie z osią N z produktem macierzowym na każdym z 2x2 matryca. W przypadku implementacji pętli wygląda ona tak, jakby:

Czy istnieje sposób, w jaki mogę to zrobić bez użycia pętli? Sprawdziłem tensordot, ale nie udało mi się go uruchomić. Myślę, że mógłbym chcieć czegoś takiego jak tensordot(a, b, axes=([1,2], [2,1])), ale to daje mi macierz NxN.

Answer 1

Wygląda na to, że wykonujesz multiplikacje macierzy dla każdego plasterka wzdłuż pierwszej osi. Za to samo, można użyć np.einsum jak tak -

np.einsum('ijk,ikl->ijl',A,B) 

testy Runtime i weryfikacji wyników -

In [179]: N = 10000 
    ...: A = np.random.rand(N,2,2) 
    ...: B = np.random.rand(N,2,2) 
    ...: 
    ...: def einsum_based(A,B): 
    ...:  return np.einsum('ijk,ikl->ijl',A,B) 
    ...: 
    ...: def forloop(A,B): 
    ...:  N = A.shape[0] 
    ...:  C = np.zeros((N,2,2)) 
    ...:  for i in range(N): 
    ...:   C[i] = np.dot(A[i], B[i]) 
    ...:  return C 
    ...: 

In [180]: np.allclose(einsum_based(A,B),forloop(A,B)) 
Out[180]: True 

In [181]: %timeit forloop(A,B) 
10 loops, best of 3: 54.9 ms per loop 

In [182]: %timeit einsum_based(A,B) 
100 loops, best of 3: 5.92 ms per loop 

Answer 2

Trzeba tylko wykonać operację na pierwszym wymiarem swoimi tensorów, która jest oznaczona przez 0:

c = tensordot(a, b, axes=(0,0)) 

To zadziała, jak chcesz. Również nie potrzebujesz listy osi, ponieważ jest to tylko jeden wymiar, w którym wykonujesz operację. Dzięki axes([1,2],[2,1]) krzyżujesz pomnożenie 2. i 3. wymiaru. Jeśli napiszesz to w notacji indeksu (konwencja podsumowująca Einsteina), odpowiada to c[i,j] = a[i,k,l]*b[j,k,l], więc zamawiasz indeksy, które chcesz zachować.

EDYCJA: Ok, problemem jest to, że tensorowy produkt z dwóch obiektów 3d to obiekt 6d. Ponieważ skurcze obejmują pary indeksów, nie ma możliwości, aby uzyskać obiekt 3D przez operację tensordot. Sztuczka polega na podzieleniu obliczeń na dwie części: najpierw wykonujesz tensordot na indeksie, aby wykonać operację macierzy, a następnie wykonujesz przekątną tensora, aby zredukować swój obiekt 4d do 3d. W jednym poleceniem:

d = np.diagonal(np.tensordot(a,b,axes=()), axis1=0, axis2=2) 

w notacji tensora d[i,j,k] = c[i,j,i,k] = a[i,j,l]*b[i,l,k].