tf.einsum

Tensor contraction over specified indices and outer product.

View aliases

Main aliases

tf.linalg.einsum

Compat aliases for migration

See Migration guide for more details.

tf.compat.v1.einsum, tf.compat.v1.linalg.einsum

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)

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	optimize: Optimization strategy to use to find contraction path using opt_einsum. Must be 'greedy', 'optimal', 'branch-2', 'branch-all' or 'auto'. (optional, default: 'greedy'). name: A name for the operation (optional).

Returns
The contracted Tensor, with shape determined by equation.

Raises
ValueError	If the format of equation is incorrect, number of inputs or their shapes are inconsistent with equation.