numpy.einsum#

numpy.einsum(subscripts,*operands,out=None,dtype=None,order='K',casting='safe',optimize=False)[source]#

Evaluates the Einstein summation convention on the operands.

Using the Einstein summation convention, many common multi-dimensional,linear algebraic array operations can be represented in a simple fashion.Inimplicit modeeinsum computes these values.

Inexplicit mode,einsum provides further flexibility to computeother array operations that might not be considered classical Einsteinsummation operations, by disabling, or forcing summation over specifiedsubscript labels.

See the notes and examples for clarification.

Parameters:
subscriptsstr

Specifies the subscripts for summation as comma separated list ofsubscript labels. An implicit (classical Einstein summation)calculation is performed unless the explicit indicator ‘->’ isincluded as well as subscript labels of the precise output form.

operandslist of array_like

These are the arrays for the operation.

outndarray, optional

If provided, the calculation is done into this array.

dtype{data-type, None}, optional

If provided, forces the calculation to use the data type specified.Note that you may have to also give a more liberalcastingparameter to allow the conversions. Default is None.

order{‘C’, ‘F’, ‘A’, ‘K’}, optional

Controls the memory layout of the output. ‘C’ means it shouldbe C contiguous. ‘F’ means it should be Fortran contiguous,‘A’ means it should be ‘F’ if the inputs are all ‘F’, ‘C’ otherwise.‘K’ means it should be as close to the layout as the inputs asis possible, including arbitrarily permuted axes.Default is ‘K’.

casting{‘no’, ‘equiv’, ‘safe’, ‘same_kind’, ‘unsafe’}, optional

Controls what kind of data casting may occur. Setting this to‘unsafe’ is not recommended, as it can adversely affect accumulations.

  • ‘no’ means the data types should not be cast at all.

  • ‘equiv’ means only byte-order changes are allowed.

  • ‘safe’ means only casts which can preserve values are allowed.

  • ‘same_kind’ means only safe casts or casts within a kind,like float64 to float32, are allowed.

  • ‘unsafe’ means any data conversions may be done.

Default is ‘safe’.

optimize{False, True, ‘greedy’, ‘optimal’}, optional

Controls if intermediate optimization should occur. No optimizationwill occur if False and True will default to the ‘greedy’ algorithm.Also accepts an explicit contraction list from thenp.einsum_pathfunction. Seenp.einsum_path for more details. Defaults to False.

Returns:
outputndarray

The calculation based on the Einstein summation convention.

See also

einsum_path,dot,inner,outer,tensordot,linalg.multi_dot
einsum

Similar verbose interface is provided by theeinops package to cover additional operations: transpose, reshape/flatten, repeat/tile, squeeze/unsqueeze and reductions. Theopt_einsum optimizes contraction order for einsum-like expressions in backend-agnostic manner.

Notes

The Einstein summation convention can be used to computemany multi-dimensional, linear algebraic array operations.einsumprovides a succinct way of representing these.

A non-exhaustive list of these operations,which can be computed byeinsum, is shown below along with examples:

The subscripts string is a comma-separated list of subscript labels,where each label refers to a dimension of the corresponding operand.Whenever a label is repeated it is summed, sonp.einsum('i,i',a,b)is equivalent tonp.inner(a,b). If a labelappears only once, it is not summed, sonp.einsum('i',a)produces a view ofa with no changes. A further examplenp.einsum('ij,jk',a,b) describes traditional matrix multiplicationand is equivalent tonp.matmul(a,b).Repeated subscript labels in one operand take the diagonal.For example,np.einsum('ii',a) is equivalent tonp.trace(a).

Inimplicit mode, the chosen subscripts are importantsince the axes of the output are reordered alphabetically. Thismeans thatnp.einsum('ij',a) doesn’t affect a 2D array, whilenp.einsum('ji',a) takes its transpose. Additionally,np.einsum('ij,jk',a,b) returns a matrix multiplication, while,np.einsum('ij,jh',a,b) returns the transpose of themultiplication since subscript ‘h’ precedes subscript ‘i’.

Inexplicit mode the output can be directly controlled byspecifying output subscript labels. This requires theidentifier ‘->’ as well as the list of output subscript labels.This feature increases the flexibility of the function sincesumming can be disabled or forced when required. The callnp.einsum('i->',a) is likenp.sum(a)ifa is a 1-D array, andnp.einsum('ii->i',a)is likenp.diag(a) ifa is a square 2-D array.The difference is thateinsum does not allow broadcasting by default.Additionallynp.einsum('ij,jh->ih',a,b) directly specifies theorder of the output subscript labels and therefore returns matrixmultiplication, unlike the example above in implicit mode.

To enable and control broadcasting, use an ellipsis. DefaultNumPy-style broadcasting is done by adding an ellipsisto the left of each term, likenp.einsum('...ii->...i',a).np.einsum('...i->...',a) is likenp.sum(a,axis=-1) for arraya of any shape.To take the trace along the first and last axes,you can donp.einsum('i...i',a), or to do a matrix-matrixproduct with the left-most indices instead of rightmost, one can donp.einsum('ij...,jk...->ik...',a,b).

When there is only one operand, no axes are summed, and no outputparameter is provided, a view into the operand is returned insteadof a new array. Thus, taking the diagonal asnp.einsum('ii->i',a)produces a view (changed in version 1.10.0).

einsum also provides an alternative way to provide the subscripts andoperands aseinsum(op0,sublist0,op1,sublist1,...,[sublistout]).If the output shape is not provided in this formateinsum will becalculated in implicit mode, otherwise it will be performed explicitly.The examples below have correspondingeinsum calls with the twoparameter methods.

Views returned from einsum are now writeable whenever the input arrayis writeable. For example,np.einsum('ijk...->kji...',a) will nowhave the same effect asnp.swapaxes(a,0,2)andnp.einsum('ii->i',a) will return a writeable view of the diagonalof a 2D array.

Added theoptimize argument which will optimize the contraction orderof an einsum expression. For a contraction with three or more operandsthis can greatly increase the computational efficiency at the cost ofa larger memory footprint during computation.

Typically a ‘greedy’ algorithm is applied which empirical tests have shownreturns the optimal path in the majority of cases. In some cases ‘optimal’will return the superlative path through a more expensive, exhaustivesearch. For iterative calculations it may be advisable to calculatethe optimal path once and reuse that path by supplying it as an argument.An example is given below.

Seenumpy.einsum_path for more details.

Examples

>>>a=np.arange(25).reshape(5,5)>>>b=np.arange(5)>>>c=np.arange(6).reshape(2,3)

Trace of a matrix:

>>>np.einsum('ii',a)60>>>np.einsum(a,[0,0])60>>>np.trace(a)60

Extract the diagonal (requires explicit form):

>>>np.einsum('ii->i',a)array([ 0,  6, 12, 18, 24])>>>np.einsum(a,[0,0],[0])array([ 0,  6, 12, 18, 24])>>>np.diag(a)array([ 0,  6, 12, 18, 24])

Sum over an axis (requires explicit form):

>>>np.einsum('ij->i',a)array([ 10,  35,  60,  85, 110])>>>np.einsum(a,[0,1],[0])array([ 10,  35,  60,  85, 110])>>>np.sum(a,axis=1)array([ 10,  35,  60,  85, 110])

For higher dimensional arrays summing a single axis can be donewith ellipsis:

>>>np.einsum('...j->...',a)array([ 10,  35,  60,  85, 110])>>>np.einsum(a,[Ellipsis,1],[Ellipsis])array([ 10,  35,  60,  85, 110])

Compute a matrix transpose, or reorder any number of axes:

>>>np.einsum('ji',c)array([[0, 3],       [1, 4],       [2, 5]])>>>np.einsum('ij->ji',c)array([[0, 3],       [1, 4],       [2, 5]])>>>np.einsum(c,[1,0])array([[0, 3],       [1, 4],       [2, 5]])>>>np.transpose(c)array([[0, 3],       [1, 4],       [2, 5]])

Vector inner products:

>>>np.einsum('i,i',b,b)30>>>np.einsum(b,[0],b,[0])30>>>np.inner(b,b)30

Matrix vector multiplication:

>>>np.einsum('ij,j',a,b)array([ 30,  80, 130, 180, 230])>>>np.einsum(a,[0,1],b,[1])array([ 30,  80, 130, 180, 230])>>>np.dot(a,b)array([ 30,  80, 130, 180, 230])>>>np.einsum('...j,j',a,b)array([ 30,  80, 130, 180, 230])

Broadcasting and scalar multiplication:

>>>np.einsum('..., ...',3,c)array([[ 0,  3,  6],       [ 9, 12, 15]])>>>np.einsum(',ij',3,c)array([[ 0,  3,  6],       [ 9, 12, 15]])>>>np.einsum(3,[Ellipsis],c,[Ellipsis])array([[ 0,  3,  6],       [ 9, 12, 15]])>>>np.multiply(3,c)array([[ 0,  3,  6],       [ 9, 12, 15]])

Vector outer product:

>>>np.einsum('i,j',np.arange(2)+1,b)array([[0, 1, 2, 3, 4],       [0, 2, 4, 6, 8]])>>>np.einsum(np.arange(2)+1,[0],b,[1])array([[0, 1, 2, 3, 4],       [0, 2, 4, 6, 8]])>>>np.outer(np.arange(2)+1,b)array([[0, 1, 2, 3, 4],       [0, 2, 4, 6, 8]])

Tensor contraction:

>>>a=np.arange(60.).reshape(3,4,5)>>>b=np.arange(24.).reshape(4,3,2)>>>np.einsum('ijk,jil->kl',a,b)array([[4400., 4730.],       [4532., 4874.],       [4664., 5018.],       [4796., 5162.],       [4928., 5306.]])>>>np.einsum(a,[0,1,2],b,[1,0,3],[2,3])array([[4400., 4730.],       [4532., 4874.],       [4664., 5018.],       [4796., 5162.],       [4928., 5306.]])>>>np.tensordot(a,b,axes=([1,0],[0,1]))array([[4400., 4730.],       [4532., 4874.],       [4664., 5018.],       [4796., 5162.],       [4928., 5306.]])

Writeable returned arrays (since version 1.10.0):

>>>a=np.zeros((3,3))>>>np.einsum('ii->i',a)[:]=1>>>aarray([[1., 0., 0.],       [0., 1., 0.],       [0., 0., 1.]])

Example of ellipsis use:

>>>a=np.arange(6).reshape((3,2))>>>b=np.arange(12).reshape((4,3))>>>np.einsum('ki,jk->ij',a,b)array([[10, 28, 46, 64],       [13, 40, 67, 94]])>>>np.einsum('ki,...k->i...',a,b)array([[10, 28, 46, 64],       [13, 40, 67, 94]])>>>np.einsum('k...,jk',a,b)array([[10, 28, 46, 64],       [13, 40, 67, 94]])

Chained array operations. For more complicated contractions, speed upsmight be achieved by repeatedly computing a ‘greedy’ path or pre-computingthe ‘optimal’ path and repeatedly applying it, using aneinsum_pathinsertion (since version 1.12.0). Performance improvements can beparticularly significant with larger arrays:

>>>a=np.ones(64).reshape(2,4,8)

Basiceinsum: ~1520ms (benchmarked on 3.1GHz Intel i5.)

>>>foriterationinrange(500):..._=np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a)

Sub-optimaleinsum (due to repeated path calculation time): ~330ms

>>>foriterationinrange(500):..._=np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a,...optimize='optimal')

Greedyeinsum (faster optimal path approximation): ~160ms

>>>foriterationinrange(500):..._=np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a,optimize='greedy')

Optimaleinsum (best usage pattern in some use cases): ~110ms

>>>path=np.einsum_path('ijk,ilm,njm,nlk,abc->',a,a,a,a,a,...optimize='optimal')[0]>>>foriterationinrange(500):..._=np.einsum('ijk,ilm,njm,nlk,abc->',a,a,a,a,a,optimize=path)
On this page