python - Entendiendo el einsum de NumPy

arrays multidimensional-array

Question

Estoy luchando por entender exactamente cómo funciona el einsum . 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 A y B

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

Answer 1

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

¿Qué hace el einsum ?

Imagine que tenemos dos matrices multidimensionales, A y B Ahora supongamos que queremos ...

• multiplique A con B de una manera particular para crear una 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.

Existe una buena posibilidad de que einsum nos ayude a hacer esto más rápido y con mayor eficiencia de memoria que las combinaciones de las funciones NumPy como multiply , sum y transpose permitirán.

¿Cómo funciona el einsum ?

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 A y B elementos 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 en A alinea los primeros ejes de las dos matrices para que se pueda transmitir la multiplicación. Las filas de la matriz de productos se suman para devolver la respuesta.

Ahora, si quisiéramos usar 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:

• A tiene un eje; lo hemos etiquetado i . Y B tiene dos ejes; hemos etiquetado el eje 0 como i y el eje 1 como j .

• Al repetir la etiqueta i en ambas matrices de entrada, le estamos diciendo al einsum que estos dos ejes deben multiplicarse juntos. En otras palabras, estamos multiplicando la matriz A con cada columna de la matriz B , tal como lo hace A[:, np.newaxis] * B

• Observe que j no 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 al einsum que einsum a lo largo de este eje. En otras palabras, estamos sumando las filas de los productos, tal como lo hace .sum(axis=1) .

Eso es básicamente todo lo que necesita 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() ).

einsum embargo, lo mejor de einsum es que no crea primero una gama temporal de productos; 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 punto usando np.einsum(''ij,jk->ik'', A, B) . Aquí hay una imagen que muestra el etiquetado de A y B y la matriz de salida que obtenemos de la función:

Puede ver que la etiqueta j se repite; esto significa que estamos multiplicando las filas de A con las columnas de B Además, la etiqueta j no está incluida en la salida; estamos sumando estos productos. Las etiquetas i y k se guardan para la salida, por lo que recuperamos una matriz 2D.

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

Sumar el eje j da el producto puntual esperado, que se muestra a la derecha.

Algunos ejercicios

Para tener una mejor idea de 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 A se 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 np.inner(A, B) , 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 las matrices 2D, C y D , siempre que los ejes sean longitudes compatibles (tanto la misma longitud como una de ellas tiene longitud 1), aquí hay algunos ejemplos:

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

  np.einsum(''ii'', C)

• La multiplicación por elementos de C y la transposición de D , C * DT , se pueden escribir:

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

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

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

Answer 2

Agarrar la idea de 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, A y 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 A mientras 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 ( j en nuestro caso), eso significa que quiere que la suma de ein suceda 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 como resultado. De lo contrario, la matriz resultante tendrá dimensiones de acuerdo con la cadena de subíndice . En nuestro ejemplo, será ik . Esto es intuitivo porque sabemos que para la multiplicación de matrices, el número de columnas en la matriz A tiene que coincidir con la cantidad de filas en la matriz B que es lo que está sucediendo aquí (es decir, codificamos este conocimiento repitiendo el char j en la cadena del subíndice )

Aquí hay algunos ejemplos más que ilustran el uso / poder de numpy.einsum() en la implementación de algunas operaciones comunes de tensor o nd-array , 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 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 arr1 * arr2 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) np.square(arr) 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 de diagonal principal) (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) 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

Tal multiplicación podría ser muy útil cuando se resuelve 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 operaciones de reshape 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

Answer 3

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,j representan filas y columnas para a . j,k para b .

Para calcular el producto y alinear el eje j , necesitamos agregar un eje a a . ( b se 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]]])

j está ausente del lado derecho, por lo que sumamos sobre j que 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]])
>>>

Answer 4

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). i es el primer dim de A , el último de C ; k el último de B , primero de C j es ''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 : es la salida final que se transpone.

Para ver más de lo que le sucede a j , cambie los subíndices C 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 dimensión k al final de A y una i al frente de B , lo que da como resultado una matriz (2,3,4).

0 + 4 + 16 = 20 , 9 + 28 + 55 = 92 , etc. Suma en j y 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]