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-dimensionalarray operations can be represented in a simple fashion. This functionprovides a way to compute such summations. The best way to understand thisfunction is to try the examples below, which show how many common NumPyfunctions can be implemented as calls toeinsum.

Parameters:

subscripts : str Specifies the subscripts for summation. operands : list of array_like These are the arrays for the operation. out : {ndarray, None}, 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. Default is False.

Returns:

output : ndarray The calculation based on the Einstein summation convention.

See also

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

Notes

New in version 1.6.0.

The subscripts string is a comma-separated list of subscript labels,where each label refers to a dimension of the corresponding operand.Repeated subscripts labels in one operand take the diagonal. For example,np.einsum('ii',a) is equivalent tonp.trace(a).

Whenever a label is repeated, it is summed, sonp.einsum('i,i',a,b)is equivalent tonp.inner(a,b). If a label appears only once,it is not summed, sonp.einsum('i',a) produces a view ofawith no changes.

The order of labels in the output is by default alphabetical. Thismeans thatnp.einsum('ij',a) doesn’t affect a 2D array, whilenp.einsum('ji',a) takes its transpose.

The output can be controlled by specifying output subscript labelsas well. This specifies the label order, and allows summing tobe disallowed or forced when desired. The callnp.einsum('i->',a)is likenp.sum(a,axis=-1), andnp.einsum('ii->i',a)is likenp.diag(a). The difference is thateinsum does notallow broadcasting by default.

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).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, you 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.

An alternative way to provide the subscripts and operands is aseinsum(op0,sublist0,op1,sublist1,...,[sublistout]). The examplesbelow have correspondingeinsum calls with the two parameter methods.

New in version 1.10.0.

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.

New in version 1.12.0.

Added theoptimize argument which will optimize the contraction orderof an einsum expression. For a contraction with three or more operands thiscan greatly increase the computational efficiency at the cost of a largermemory footprint during computation.

Seenp.einsum_path for more details.

Examples

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

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

>>>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])

>>>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])

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

>>>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]])

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

>>>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]])

>>>np.einsum('i...->...',a)array([50, 55, 60, 65, 70])>>>np.einsum(a,[0,Ellipsis],[Ellipsis])array([50, 55, 60, 65, 70])>>>np.sum(a,axis=0)array([50, 55, 60, 65, 70])

>>>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.]])

>>>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]])

>>># 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.]])