tf.einsum

TensorFlow 1 version

View source on GitHub

Tensor contraction over specified indices and outer product.

tf.einsum(
    equation, *inputs, **kwargs
)

This function returns a tensor whose elements are defined by equation, which is written in a shorthand form inspired by the Einstein summation convention. 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, the equation is obtained from the more familiar element-wise equation by

1. removing variable names, brackets, and commas,
2. replacing "*" with ",",
3. dropping summation signs, and
4. moving the output to the right, and replacing "=" with "->".

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 NumPy-style broadcasting across the batch dimensions, use:

einsum('...ij,...jk->...ik', u, v)