View source on GitHub 
Tensor contraction over specified indices and outer product.
tf.einsum( equation, *inputs, **kwargs )
Einsum allows defining Tensors by defining their elementwise computation. This computation is defined by equation
, a shorthand form based on Einstein summation. As an example, consider multiplying two matrices A and B to form a matrix C. The elements of C are given by:
C[i,k] = sum_j A[i,j] * B[j,k]
The corresponding equation
is:
ij,jk>ik
In general, to convert the elementwise equation into the equation
string, use the following procedure (intermediate strings for matrix multiplication example provided in parentheses):
ik = sum_j ij * jk
)ik = sum_j ij , jk
)ik = ij, jk
)ij,jk>ik
)Many common operations can be expressed in this way. For example:
# Matrix multiplication einsum('ij,jk>ik', m0, m1) # output[i,k] = sum_j m0[i,j] * m1[j, k] # Dot product einsum('i,i>', u, v) # output = sum_i u[i]*v[i] # Outer product einsum('i,j>ij', u, v) # output[i,j] = u[i]*v[j] # Transpose einsum('ij>ji', m) # output[j,i] = m[i,j] # Trace einsum('ii', m) # output[j,i] = trace(m) = sum_i m[i, i] # Batch matrix multiplication einsum('aij,ajk>aik', s, t) # out[a,i,k] = sum_j s[a,i,j] * t[a, j, k]
To enable and control broadcasting, use an ellipsis. For example, to perform batch matrix multiplication with NumPystyle broadcasting across the batch dimensions, use:
einsum('...ij,...jk>...ik', u, v)
Args  

equation  a str describing the contraction, in the same format as numpy.einsum . 
*inputs  the inputs to contract (each one a Tensor ), whose shapes should be consistent with equation . 
**kwargs 

Returns  

The contracted Tensor , with shape determined by equation . 
Raises  

ValueError  If

© 2020 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 3.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/versions/r2.3/api_docs/python/tf/einsum