← retour aux snippets

numpy: einsum vs dot

Utiliser einsum pour des contractions génériques.

python numpy #numpy#einsum#dot

objectif

Utiliser einsum pour des contractions génériques.

code minimal

import numpy as np
A = np.arange(6).reshape(2,3)
B = np.arange(6).reshape(3,2)
d1 = A @ B
d2 = np.einsum("ik,kj->ij", A, B)
print(np.allclose(d1, d2))

utilisation

import numpy as np
x = np.arange(3); y = np.arange(3)
print(np.einsum("i,i->", x, y).item())

variante(s) utile(s)

import numpy as np
T = np.ones((2,3,4))
print(np.einsum("ijk->i", T).shape[0])

notes

  • einsum réduit allocations intermédiaires complexes.