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 multidimensional, linear algebraic array operations can be represented in a simple fashion. In implicit mode einsum
computes these values.
In explicit mode, einsum
provides further flexibility to compute other array operations that might not be considered classical Einstein summation operations, by disabling, or forcing summation over specified subscript labels.
See the notes and examples for clarification.
Parameters: 


Returns: 

See also
New in version 1.6.0.
The Einstein summation convention can be used to compute many multidimensional, linear algebraic array operations. einsum
provides a succinct way of representing these.
A nonexhaustive list of these operations, which can be computed by einsum
, is shown below along with examples:
numpy.trace
.numpy.diag
.numpy.sum
.numpy.transpose
.numpy.matmul
numpy.dot
.numpy.inner
numpy.outer
.numpy.multiply
.numpy.tensordot
.numpy.einsum_path
.The subscripts string is a commaseparated list of subscript labels, where each label refers to a dimension of the corresponding operand. Whenever a label is repeated it is summed, so np.einsum('i,i', a, b)
is equivalent to np.inner(a,b)
. If a label appears only once, it is not summed, so np.einsum('i', a)
produces a view of a
with no changes. A further example np.einsum('ij,jk', a, b)
describes traditional matrix multiplication and is equivalent to np.matmul(a,b)
. Repeated subscript labels in one operand take the diagonal. For example, np.einsum('ii', a)
is equivalent to np.trace(a)
.
In implicit mode, the chosen subscripts are important since the axes of the output are reordered alphabetically. This means that np.einsum('ij', a)
doesn’t affect a 2D array, while np.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 the multiplication since subscript ‘h’ precedes subscript ‘i’.
In explicit mode the output can be directly controlled by specifying output subscript labels. This requires the identifier ‘>’ as well as the list of output subscript labels. This feature increases the flexibility of the function since summing can be disabled or forced when required. The call np.einsum('i>', a)
is like np.sum(a, axis=1)
, and np.einsum('ii>i', a)
is like np.diag(a)
. The difference is that einsum
does not allow broadcasting by default. Additionally np.einsum('ij,jh>ih', a, b)
directly specifies the order of the output subscript labels and therefore returns matrix multiplication, unlike the example above in implicit mode.
To enable and control broadcasting, use an ellipsis. Default NumPystyle broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii>...i', a)
. To take the trace along the first and last axes, you can do np.einsum('i...i', a)
, or to do a matrixmatrix product with the leftmost indices instead of rightmost, one can do np.einsum('ij...,jk...>ik...', a, b)
.
When there is only one operand, no axes are summed, and no output parameter is provided, a view into the operand is returned instead of a new array. Thus, taking the diagonal as np.einsum('ii>i', a)
produces a view (changed in version 1.10.0).
einsum
also provides an alternative way to provide the subscripts and operands as einsum(op0, sublist0, op1, sublist1, ..., [sublistout])
. If the output shape is not provided in this format einsum
will be calculated in implicit mode, otherwise it will be performed explicitly. The examples below have corresponding einsum
calls with the two parameter methods.
New in version 1.10.0.
Views returned from einsum are now writeable whenever the input array is writeable. For example, np.einsum('ijk...>kji...', a)
will now have the same effect as np.swapaxes(a, 0, 2)
and np.einsum('ii>i', a)
will return a writeable view of the diagonal of a 2D array.
New in version 1.12.0.
Added the optimize
argument which will optimize the contraction order of an einsum expression. For a contraction with three or more operands this can greatly increase the computational efficiency at the cost of a larger memory footprint during computation.
Typically a ‘greedy’ algorithm is applied which empirical tests have shown returns the optimal path in the majority of cases. In some cases ‘optimal’ will return the superlative path through a more expensive, exhaustive search. For iterative calculations it may be advisable to calculate the optimal path once and reuse that path by supplying it as an argument. An example is given below.
See numpy.einsum_path
for more details.
>>> 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 done with 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 >>> a array([[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 ups might be achieved by repeatedly computing a ‘greedy’ path or precomputing the ‘optimal’ path and repeatedly applying it, using an einsum_path
insertion (since version 1.12.0). Performance improvements can be particularly significant with larger arrays:
>>> a = np.ones(64).reshape(2,4,8)
Basic einsum
: ~1520ms (benchmarked on 3.1GHz Intel i5.)
>>> for iteration in range(500): ... _ = np.einsum('ijk,ilm,njm,nlk,abc>',a,a,a,a,a)
Suboptimal einsum
(due to repeated path calculation time): ~330ms
>>> for iteration in range(500): ... _ = np.einsum('ijk,ilm,njm,nlk,abc>',a,a,a,a,a, optimize='optimal')
Greedy einsum
(faster optimal path approximation): ~160ms
>>> for iteration in range(500): ... _ = np.einsum('ijk,ilm,njm,nlk,abc>',a,a,a,a,a, optimize='greedy')
Optimal einsum
(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] >>> for iteration in range(500): ... _ = np.einsum('ijk,ilm,njm,nlk,abc>',a,a,a,a,a, optimize=path)
© 2005–2019 NumPy Developers
Licensed under the 3clause BSD License.
https://docs.scipy.org/doc/numpy1.17.0/reference/generated/numpy.einsum.html