Entendiendo el einsum de NumPy

¿Cómo np.einsumfunciona?

Dadas las matrices Ay B, su multiplicación de matrices seguida de transposición se calcula usando (A @ B).T, o de manera equivalente, usando:

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

Sep 29 '14 04:09 Lance Strait

Captar la idea de numpy.einsum()es muy fácil si la 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 tienes dos matrices 2D Ay By quieres realizar una multiplicación de matrices. Tu también:

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

Aquí la cadena de subíndice ij corresponde a una matriz, Amientras que la cadena de subíndice jk corresponde a una matriz B. Además, lo más importante a tener en cuenta aquí es que la cantidad 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índices ( jen nuestro caso), eso significa que desea que la einsuma se realice a lo largo de esas dimensiones. Por lo tanto, se reducirán en suma (es decir, esa dimensión desaparecerá ) .

La cadena de subíndice después de este ->símbolo representa las dimensiones de nuestra matriz resultante. Si lo deja vacío, todo se sumará y se devolverá un valor escalar como resultado. De lo contrario, la matriz resultante tendrá dimensiones según la cadena del subíndice . En nuestro ejemplo, será ik. Esto es intuitivo porque sabemos que para que la multiplicación de matrices funcione, el número de columnas en la matriz Atiene que coincidir con el número de filas en la matriz B, que 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 np.einsum()implementar algunas operaciones comunes de tensor o matriz nd , de manera sucinta.

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 de matrices (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 de Hadamard (es decir, producto por 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) Elevación al cuadrado 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) Traza (es decir, suma de elementos de la diagonal principal) (similar a np.trace(arr))

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

6) Transpuesta 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 exterior (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 matrices por 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 de 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) Productos de doble punto (similar a np.sum(hadamard-product) 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

Esta multiplicación podría resultar muy útil al resolver un sistema lineal de ecuaciones ( Ax = b ) donde desee 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 tenemos 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

Bonificación : lea más matemáticas aquí: Suma de Einstein y definitivamente aquí: Notación tensorial

Dec 25 '2017 07:12 kmario23

Al leer ecuaciones einsum, me ha parecido más útil poder reducirlas mentalmente a sus versiones imperativas.

Comencemos con la siguiente (imponente) declaración:

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

Primero, analizando la puntuación, vemos que tenemos dos manchas de 4 letras separadas por comas: bhwiy bhwj, antes de la flecha, y una sola mancha 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, dejemos que cada letra de cada blob sea el nombre de una variable de rango. La posición en la que aparece la letra en la mancha es el índice del eje que recorre en ese tensor. La suma imperativa que produce cada elemento de C, por lo tanto, tiene que comenzar con tres bucles for 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, tienes un forbucle para cada índice de salida de C. Dejaremos los rangos indeterminados por ahora.

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

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 bucle más interno ahora tenemos todos los índices definidos, por lo que podemos escribir la suma real y la traducción estará 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 has podido seguir el código hasta ahora, ¡felicidades! Esto es todo lo que necesitas para poder leer ecuaciones einsum. Observe en particular cómo la fórmula einsum original se relaciona con la declaración de suma final en el fragmento anterior. Los bucles for y los límites de rango son simplemente tontos y esa declaración final es todo lo que realmente necesitas para entender lo que está pasando.

Para completar, 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 tienen que 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]

Jan 22 '2020 11:01 Stefan Dragnev