Entendiendo el einsum de NumPy

Estoy luchando por entender exactamente cómo einsumfunciona. He visto la documentación y algunos ejemplos, pero no parece que se quede.

Aquí hay un ejemplo que analizamos en clase:

C = np.einsum("ij,jk->ki", A, B)

para dos matrices AyB

Creo que esto tomaría A^T * B, pero no estoy seguro (está tomando la transposición de uno de ellos, ¿verdad?). ¿Alguien puede explicarme exactamente lo que está sucediendo aquí (y en general cuando se usa einsum)?

(Nota: esta respuesta se basa en una breve publicación de blog sobre lo einsumque escribí hace un tiempo).

¿Qué einsumhacer?

Imagine que tenemos dos matrices multidimensionales Ay B. Ahora supongamos que queremos ...

• multiplican A con Bde una manera particular para crear nueva gama de productos; y luego tal vez
• suma esta nueva matriz a lo largo de ejes particulares; y luego tal vez
• transponer los ejes de la nueva matriz en un orden particular.

Hay una posibilidad buena que einsumnos ayudará a hacer esto más rápido y más memoria-eficiente que las combinaciones de las funciones NumPy gusta multiply, sumy transposelo permita.

Como einsumfunciona

Aquí hay un ejemplo simple (pero no completamente trivial). Tome las siguientes dos matrices:

A = np.array([0, 1, 2])

B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

Multiplicaremos Ay en función de los Belementos y luego sumaremos a lo largo de las filas de la nueva matriz. En NumPy "normal" escribiríamos:

>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])

Entonces, aquí, la operación de indexación Aalinea los primeros ejes de las dos matrices para que la multiplicación pueda transmitirse. Las filas de la matriz de productos se suman para devolver la respuesta.

Ahora, si quisiéramos usarlo einsum, podríamos escribir:

>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])

La cadena de firma'i,ij->i' es la clave aquí y necesita un poco de explicación. Puedes pensarlo en dos mitades. En el lado izquierdo (a la izquierda de ->) hemos etiquetado las dos matrices de entrada. A la derecha de ->, hemos etiquetado la matriz con la que queremos terminar.

Esto es lo que sucede después:

• Atiene un eje; Lo hemos etiquetado i. Y Btiene dos ejes; hemos etiquetado el eje 0 como iy el eje 1 como j.

• Al repetir la etiqueta ien ambas matrices de entrada, estamos diciendo einsumque estos dos ejes deben multiplicarse juntos. En otras palabras, estamos multiplicando la matriz Acon cada columna de la matriz B, al igual que lo A[:, np.newaxis] * Bhace.

• Tenga en cuenta que jno aparece como una etiqueta en nuestra salida deseada; acabamos de usar i(queremos terminar con una matriz 1D). Al omitir la etiqueta, le estamos diciendo einsumque sume a lo largo de este eje. En otras palabras, estamos sumando las filas de los productos, al igual que lo .sum(axis=1)hace.

Eso es básicamente todo lo que necesitas saber para usar einsum. Ayuda a jugar un poco; Si dejamos ambas etiquetas en la salida, 'i,ij->ij'recuperamos una matriz 2D de productos (igual que A[:, np.newaxis] * B). Si decimos que no hay etiquetas de salida, 'i,ij->recuperamos un solo número (lo mismo que hacer (A[:, np.newaxis] * B).sum()).

La gran cosa acerca einsumsin embargo, es que es no construir una matriz temporal de productos de primera; simplemente resume los productos a medida que avanza. Esto puede conducir a grandes ahorros en el uso de la memoria.

Un ejemplo un poco más grande

Para explicar el producto punto, aquí hay dos nuevas matrices:

A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])

B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

Calcularemos el producto de punto usando np.einsum('ij,jk->ik', A, B). Aquí hay una foto que muestra el etiquetado de la Ae By la matriz de salida que obtenemos de la función:

Puede ver que la etiqueta jse repite, esto significa que estamos multiplicando las filas de Acon las columnas de B. Además, la etiqueta jno está incluida en la salida; estamos sumando estos productos. Las etiquetas iy kse guardan para la salida, por lo que recuperamos una matriz 2D.

Puede ser que sea aún más clara para comparar este resultado con la matriz donde la etiqueta jse no se suma. A continuación, a la izquierda, puede ver la matriz 3D que resulta de la escritura np.einsum('ij,jk->ijk', A, B)(es decir, hemos mantenido la etiqueta j):

El eje de suma jda el producto de punto esperado, que se muestra a la derecha.

Algunos ejercicios

Para tener más experiencia einsum, puede ser útil implementar operaciones familiares de matriz NumPy utilizando la notación de subíndice. Cualquier cosa que implique combinaciones de ejes de multiplicación y suma se puede escribir usando einsum.

Deje que A y B sean dos matrices 1D con la misma longitud. Por ejemplo, A = np.arange(10)y B = np.arange(5, 15).

• La suma de Ase puede escribir:

  np.einsum('i->', A)

• La multiplicación por elementos A * B, se puede escribir:

  np.einsum('i,i->i', A, B)

• El producto interno o producto de punto, np.inner(A, B)o np.dot(A, B), se puede escribir:

  np.einsum('i,i->', A, B) # or just use 'i,i'

• El producto externo np.outer(A, B), se puede escribir:

  np.einsum('i,j->ij', A, B)

Para matrices 2D Cy D, siempre que los ejes sean longitudes compatibles (tanto la misma longitud como una de ellas tiene longitud 1), aquí hay algunos ejemplos:

• El rastro de C(suma de la diagonal principal) np.trace(C), se puede escribir:

  np.einsum('ii', C)

• En cuanto al elemento de multiplicación de Cy la transpuesta de D, C * D.T, se puede escribir:

  np.einsum('ij,ji->ij', C, D)

• Multiplicando cada elemento de Cpor la matriz D(para hacer una matriz 4D) C[:, :, None, None] * D, se puede escribir:

  np.einsum('ij,kl->ijkl', C, D)  

Captar la idea numpy.einsum()es muy fácil si lo entiendes intuitivamente. Como ejemplo, comencemos con una descripción simple que involucra la multiplicación de matrices .

Para usar numpy.einsum(), todo lo que tiene que hacer es pasar la llamada cadena de subíndices como argumento, seguida de sus matrices de entrada .

Digamos que tiene dos matrices 2D, Ay B, y desea hacer una multiplicación de matrices. Tu también:

np.einsum("ij, jk -> ik", A, B)

Aquí la cadena de subíndice ij corresponde a la matriz, Amientras que la cadena de subíndice jk corresponde a la matriz B. Además, lo más importante a tener en cuenta aquí es que el número de caracteres en cada cadena de subíndice debe coincidir con las dimensiones de la matriz. (es decir, dos caracteres para matrices 2D, tres caracteres para matrices 3D, etc.) Y si repite los caracteres entre cadenas de subíndice ( jen nuestro caso), eso significa que desea que la einsuma ocurra a lo largo de esas dimensiones. Por lo tanto, se reducirán la suma. (es decir, esa dimensión se habrá ido )

La cadena de subíndice después de esto ->, será nuestra matriz resultante. Si lo deja vacío, todo se sumará y se devolverá un valor escalar. De lo contrario, la matriz resultante tendrá dimensiones de acuerdo con la cadena de subíndice . En nuestro ejemplo, lo será ik. Esto es intuitivo porque sabemos que para la multiplicación de matrices, el número de columnas en la matriz Adebe coincidir con el número de filas en la matriz, Bque es lo que está sucediendo aquí (es decir, codificamos este conocimiento repitiendo el carácter jen la cadena del subíndice )

Aquí hay algunos ejemplos más que ilustran el uso / poder de la np.einsum()implementación de algunas operaciones comunes de tensor o nd-array , sucintamente.

Entradas

# a vector
In [197]: vec
Out[197]: array([0, 1, 2, 3])

# an array
In [198]: A
Out[198]: 
array([[11, 12, 13, 14],
       [21, 22, 23, 24],
       [31, 32, 33, 34],
       [41, 42, 43, 44]])

# another array
In [199]: B
Out[199]: 
array([[1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3],
       [4, 4, 4, 4]])

1) Multiplicación matricial (similar a np.matmul(arr1, arr2))

In [200]: np.einsum("ij, jk -> ik", A, B)
Out[200]: 
array([[130, 130, 130, 130],
       [230, 230, 230, 230],
       [330, 330, 330, 330],
       [430, 430, 430, 430]])

2) Extraer elementos a lo largo de la diagonal principal (similar a np.diag(arr))

In [202]: np.einsum("ii -> i", A)
Out[202]: array([11, 22, 33, 44])

3) Producto Hadamard (es decir, producto basado en elementos de dos matrices) (similar a arr1 * arr2)

In [203]: np.einsum("ij, ij -> ij", A, B)
Out[203]: 
array([[ 11,  12,  13,  14],
       [ 42,  44,  46,  48],
       [ 93,  96,  99, 102],
       [164, 168, 172, 176]])

4) Cuadratura por elementos (similar a np.square(arr)o arr ** 2)

In [210]: np.einsum("ij, ij -> ij", B, B)
Out[210]: 
array([[ 1,  1,  1,  1],
       [ 4,  4,  4,  4],
       [ 9,  9,  9,  9],
       [16, 16, 16, 16]])

5) Trace (es decir, suma de elementos principal-diagonales) (similar a np.trace(arr))

In [217]: np.einsum("ii -> ", A)
Out[217]: 110

6) Transposición de matriz (similar a np.transpose(arr))

In [221]: np.einsum("ij -> ji", A)
Out[221]: 
array([[11, 21, 31, 41],
       [12, 22, 32, 42],
       [13, 23, 33, 43],
       [14, 24, 34, 44]])

7) Producto externo (de vectores) (similar a np.outer(vec1, vec2))

In [255]: np.einsum("i, j -> ij", vec, vec)
Out[255]: 
array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6],
       [0, 3, 6, 9]])

8) Producto interno (de vectores) (similar a np.inner(vec1, vec2))

In [256]: np.einsum("i, i -> ", vec, vec)
Out[256]: 14

9) Suma a lo largo del eje 0 (similar a np.sum(arr, axis=0))

In [260]: np.einsum("ij -> j", B)
Out[260]: array([10, 10, 10, 10])

10) Suma a lo largo del eje 1 (similar a np.sum(arr, axis=1))

In [261]: np.einsum("ij -> i", B)
Out[261]: array([ 4,  8, 12, 16])

11) Multiplicación de matriz de lotes

In [287]: BM = np.stack((A, B), axis=0)

In [288]: BM
Out[288]: 
array([[[11, 12, 13, 14],
        [21, 22, 23, 24],
        [31, 32, 33, 34],
        [41, 42, 43, 44]],

       [[ 1,  1,  1,  1],
        [ 2,  2,  2,  2],
        [ 3,  3,  3,  3],
        [ 4,  4,  4,  4]]])

In [289]: BM.shape
Out[289]: (2, 4, 4)

# batch matrix multiply using einsum
In [292]: BMM = np.einsum("bij, bjk -> bik", BM, BM)

In [293]: BMM
Out[293]: 
array([[[1350, 1400, 1450, 1500],
        [2390, 2480, 2570, 2660],
        [3430, 3560, 3690, 3820],
        [4470, 4640, 4810, 4980]],

       [[  10,   10,   10,   10],
        [  20,   20,   20,   20],
        [  30,   30,   30,   30],
        [  40,   40,   40,   40]]])

In [294]: BMM.shape
Out[294]: (2, 4, 4)

12) Suma a lo largo del eje 2 (similar a np.sum(arr, axis=2))

In [330]: np.einsum("ijk -> ij", BM)
Out[330]: 
array([[ 50,  90, 130, 170],
       [  4,   8,  12,  16]])

13) Suma todos los elementos en la matriz (similar a np.sum(arr))

In [335]: np.einsum("ijk -> ", BM)
Out[335]: 480

14) Suma sobre múltiples ejes (es decir, marginación)
(similar a np.sum(arr, axis=(axis0, axis1, axis2, axis3, axis4, axis6, axis7)))

# 8D array
In [354]: R = np.random.standard_normal((3,5,4,6,8,2,7,9))

# marginalize out axis 5 (i.e. "n" here)
In [363]: esum = np.einsum("ijklmnop -> n", R)

# marginalize out axis 5 (i.e. sum over rest of the axes)
In [364]: nsum = np.sum(R, axis=(0,1,2,3,4,6,7))

In [365]: np.allclose(esum, nsum)
Out[365]: True

15) de doble punto los productos (similar a np.sum (Hadamard-producto) cf. 3 )

In [772]: A
Out[772]: 
array([[1, 2, 3],
       [4, 2, 2],
       [2, 3, 4]])

In [773]: B
Out[773]: 
array([[1, 4, 7],
       [2, 5, 8],
       [3, 6, 9]])

In [774]: np.einsum("ij, ij -> ", A, B)
Out[774]: 124

16) Multiplicación de matrices 2D y 3D

Tal multiplicación podría ser muy útil al resolver un sistema lineal de ecuaciones ( Ax = b ) donde desea verificar el resultado.

# inputs
In [115]: A = np.random.rand(3,3)
In [116]: b = np.random.rand(3, 4, 5)

# solve for x
In [117]: x = np.linalg.solve(A, b.reshape(b.shape[0], -1)).reshape(b.shape)

# 2D and 3D array multiplication :)
In [118]: Ax = np.einsum('ij, jkl', A, x)

# indeed the same!
In [119]: np.allclose(Ax, b)
Out[119]: True

Por el contrario, si uno tiene que usar np.matmul()para esta verificación, tenemos que hacer un par de reshapeoperaciones para lograr el mismo resultado como:

# reshape 3D array `x` to 2D, perform matmul
# then reshape the resultant array to 3D
In [123]: Ax_matmul = np.matmul(A, x.reshape(x.shape[0], -1)).reshape(x.shape)

# indeed correct!
In [124]: np.allclose(Ax, Ax_matmul)
Out[124]: True

Bono : Lea más matemáticas aquí: Einstein-Summation y definitivamente aquí: Tensor-Notation

Hagamos 2 matrices, con dimensiones diferentes pero compatibles para resaltar su interacción

In [43]: A=np.arange(6).reshape(2,3)
Out[43]: 
array([[0, 1, 2],
       [3, 4, 5]])


In [44]: B=np.arange(12).reshape(3,4)
Out[44]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

Su cálculo, toma un 'punto' (suma de productos) de un (2,3) con un (3,4) para producir una matriz (4,2). ies el primer dim de A, el último de C; kel último de B, primero de C. jes 'consumido' por la sumatoria.

In [45]: C=np.einsum('ij,jk->ki',A,B)
Out[45]: 
array([[20, 56],
       [23, 68],
       [26, 80],
       [29, 92]])

Esto es lo mismo que np.dot(A,B).T: se transpone la salida final.

Para ver más de lo que sucede j, cambie los Csubíndices a ijk:

In [46]: np.einsum('ij,jk->ijk',A,B)
Out[46]: 
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]]])

Esto también se puede producir con:

A[:,:,None]*B[None,:,:]

Es decir, agregue una kdimensión al final de A, y una ial frente de B, dando como resultado una matriz (2,3,4).

0 + 4 + 16 = 20, 9 + 28 + 55 = 92, Etc; Suma jy transpone para obtener el resultado anterior:

np.sum(A[:,:,None] * B[None,:,:], axis=1).T

# C[k,i] = sum(j) A[i,j (,k) ] * B[(i,)  j,k]

Encontré NumPy: Los trucos del oficio (Parte II) instructivos

Usamos -> para indicar el orden de la matriz de salida. Así que piense que 'ij, i-> j' tiene el lado izquierdo (LHS) y el lado derecho (RHS). Cualquier repetición de etiquetas en el LHS calcula el elemento del producto sabiamente y luego suma. Al cambiar la etiqueta en el lado RHS (salida), podemos definir el eje en el que queremos proceder con respecto a la matriz de entrada, es decir, la suma a lo largo del eje 0, 1 y así sucesivamente.

import numpy as np

>>> a
array([[1, 1, 1],
       [2, 2, 2],
       [3, 3, 3]])
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> d = np.einsum('ij, jk->ki', a, b)

Observe que hay tres ejes, i, j, k, y que j se repite (en el lado izquierdo). i,jrepresentar filas y columnas para a. j,kpara b.

Para calcular el producto y alinear el jeje, necesitamos agregarle un eje a. ( bse transmitirá a lo largo (?) del primer eje)

a[i, j, k]
   b[j, k]

>>> c = a[:,:,np.newaxis] * b
>>> c
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 0,  2,  4],
        [ 6,  8, 10],
        [12, 14, 16]],

       [[ 0,  3,  6],
        [ 9, 12, 15],
        [18, 21, 24]]])

jestá ausente del lado derecho, por lo que sumamos jcuál es el segundo eje de la matriz 3x3x3

>>> c = c.sum(1)
>>> c
array([[ 9, 12, 15],
       [18, 24, 30],
       [27, 36, 45]])

Finalmente, los índices se invierten (alfabéticamente) en el lado derecho, por lo que transponemos.

>>> c.T
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])

>>> np.einsum('ij, jk->ki', a, b)
array([[ 9, 18, 27],
       [12, 24, 36],
       [15, 30, 45]])
>>>

Al leer las ecuaciones de einsum, me pareció más útil poder reducirlas mentalmente a sus versiones imperativas.

Comencemos con la siguiente declaración (imponente):

C = np.einsum('bhwi,bhwj->bij', A, B)

Al trabajar primero en la puntuación, vemos que tenemos dos blobs separados por comas de 4 letras, bhwiy bhwj, antes de la flecha, y un único blob de 3 letras bijdespués. Por lo tanto, la ecuación produce un resultado de tensor de rango 3 a partir de dos entradas de tensor de rango 4.

Ahora, deje que cada letra de cada blob sea el nombre de una variable de rango. La posición en la que aparece la letra en el blob es el índice del eje sobre el que se extiende en ese tensor. La suma imperativa que produce cada elemento de C, por lo tanto, tiene que comenzar con tres bucles anidados, uno para cada índice de C.

for b in range(...):
    for i in range(...):
        for j in range(...):
            # the variables b, i and j index C in the order of their appearance in the equation
            C[b, i, j] = ...

Entonces, esencialmente, tiene un forciclo para cada índice de salida de C. Dejaremos los rangos indeterminados por ahora.

A continuación, miramos el lado izquierdo: ¿hay alguna variable de rango allí que no aparezca en el lado derecho ? En nuestro caso, sí, hy w. Agregue un forbucle anidado interno para cada variable:

for b in range(...):
    for i in range(...):
        for j in range(...):
            C[b, i, j] = 0
            for h in range(...):
                for w in range(...):
                    ...

Dentro del ciclo más interno ahora tenemos todos los índices definidos, por lo que podemos escribir la suma real y la traducción está completa:

# three nested for-loops that index the elements of C
for b in range(...):
    for i in range(...):
        for j in range(...):

            # prepare to sum
            C[b, i, j] = 0

            # two nested for-loops for the two indexes that don't appear on the right-hand side
            for h in range(...):
                for w in range(...):
                    # Sum! Compare the statement below with the original einsum formula
                    # 'bhwi,bhwj->bij'

                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]

Si hasta ahora ha podido seguir el código, ¡felicidades! Esto es todo lo que necesitas para poder leer las ecuaciones de einsum. Observe en particular cómo la fórmula original de einsum se asigna a la declaración de suma final en el fragmento de arriba. Los for-loops y los límites de rango son simplemente imprecisos y esa declaración final es todo lo que realmente necesita para comprender lo que está sucediendo.

En aras de la exhaustividad, veamos cómo determinar los rangos para cada variable de rango. Bueno, el rango de cada variable es simplemente la longitud de las dimensiones que indexa. Obviamente, si una variable indexa más de una dimensión en uno o más tensores, entonces las longitudes de cada una de esas dimensiones deben ser iguales. Aquí está el código anterior con los rangos completos:

# C's shape is determined by the shapes of the inputs
# b indexes both A and B, so its range can come from either A.shape or B.shape
# i indexes only A, so its range can only come from A.shape, the same is true for j and B
assert A.shape[0] == B.shape[0]
assert A.shape[1] == B.shape[1]
assert A.shape[2] == B.shape[2]
C = np.zeros((A.shape[0], A.shape[3], B.shape[3]))
for b in range(A.shape[0]): # b indexes both A and B, or B.shape[0], which must be the same
    for i in range(A.shape[3]):
        for j in range(B.shape[3]):
            # h and w can come from either A or B
            for h in range(A.shape[1]):
                for w in range(A.shape[2]):
                    C[b, i, j] += A[b, h, w, i] * B[b, h, w, j]