Skip to content

einsum

Evaluate an Einstein summation equation on symmetry-aware tensors.

einsum

einsum(equation: str, *tensors: Tensor) -> Tensor

Evaluate an Einstein summation equation on symmetry-aware tensors.

Parses equation and dispatches to contract, trace, and permute. Three equation types are supported:

  1. Permutation — single tensor, output is a reordering of the input subscript:

    einsum('ij->ji', A)
    
  2. Trace — single tensor with a repeated subscript letter. The output subscript lists the surviving axes in the desired order:

    einsum('ii->', A)       # full trace → scalar
    einsum('iijk->jk', A)   # partial trace
    
  3. Sequential contraction — two or more tensors contracted from left to right. At each step, letters present in both the current result and the next tensor are contracted, unless they appear in the output subscript:

    einsum('ij,jk->ik', A, B)       # matrix multiply
    einsum('ij,jk,kl->il', A, B, C) # chain contraction
    

Parameters:

Name Type Description Default
equation str

Subscript equation of the form '<lhs>-><rhs>'. <lhs> is a comma-separated list of subscript strings, one per input tensor. <rhs> is the output subscript. The -> separator is required; implicit output subscripts are not supported.

Each subscript character must be a single ASCII letter. A letter may appear at most twice within a single input subscript (forming one evaluation pair).

required
*tensors Tensor

Input Tensor objects. The number of tensors must match the number of comma-separated subscript strings in equation, and the order of each tensor must equal the length of its subscript.

()

Returns:

Type Description
Tensor

Result tensor. For a full trace this is a scalar (order 0). For a partial trace or contraction the result has the axes listed in the output subscript in that order.

Raises:

Type Description
ValueError

Raised in any of the following situations:

  • equation does not contain '->'.
  • The number of comma-separated input subscripts does not match the number of tensors supplied.
  • A subscript length does not match the order of the corresponding tensor.
  • A letter appears three or more times within a single input subscript.
  • An error is propagated from trace, e.g. a trace pair whose two axes share the same direction.
  • An error is propagated from contract, e.g. contracted axes have mismatched itags or the same direction.
Notes

No contraction-order optimisation. For multi-tensor equations the tensors are contracted strictly from left to right. This may be suboptimal for networks where a different pairing would reduce intermediate tensor sizes.

Vectors (1st order tensors) are not supported by this library. Any equation whose intermediate or final result would have exactly one index will raise an error from the underlying contract or trace call.

Hadamard (batch) indices — a letter that appears in two or more input subscripts and also in the output subscript — are not supported. Such a letter would need to survive as a free axis in both tensors without being summed over, which requires element-wise multiplication semantics that contract does not provide.

Examples:

Permutation:

>>> einsum('ij->ji', A)

Trace to scalar:

>>> einsum('ii->', A)

Matrix multiplication:

>>> einsum('ij,jk->ik', A, B)

Chain contraction:

>>> einsum('ij,jk,kl->il', A, B, C)

Description

einsum parses a subscript equation string and dispatches to contract, trace, and permute to carry out the requested operation. Three core equation types are supported — permutation, trace, and contraction — and they can be freely combined. Some representative patterns:

  1. Permutation — a single input tensor whose output subscript is a reordering of the input subscript:

    einsum('ij->ji', A)           # transpose
    einsum('abcd->dcba', A)       # reverse all axes
    
  2. Trace — a repeated letter within one input subscript causes those two axes to be traced out. Remaining axes are kept in the order given by the output subscript:

    einsum('ii->', A)             # full trace → scalar
    einsum('iijk->jk', A)         # partial trace, axes 0 and 1
    einsum('abiicd->abcd', A)     # trace middle axes 2 and 3
    
  3. Contraction — a letter shared between two input subscripts but absent from the output subscript is summed over at that step. Tensors are contracted strictly from left to right:

    einsum('ij,jk->ik', A, B)         # matrix multiply
    einsum('ij,jk->ki', A, B)         # matrix multiply then transpose
    einsum('ij,ji->', A, B)           # full contraction to scalar
    einsum('ijk,kjl->il', A, B)       # multi-index, non-trivial axis order
    einsum('ij,jk,kl->il', A, B, C)  # chain contraction
    
  4. Outer product — no shared letters between two inputs means no summation; the result carries all axes from both tensors:

    einsum('ij,kl->ijkl', A, B)       # outer product, natural order
    einsum('ij,kl->klij', A, B)       # outer product with output permutation
    
  5. Mixed trace and contraction — within-tensor repeated letters are traced first; the result then participates in the contraction chain as normal:

    einsum('iijk,kl->jl', A, B)       # trace A on axes 0,1 then contract on k
    einsum('ijk,kjl,lmn->imn', A, B, C)  # non-trivial chain with trace step
    

Equation syntax

An equation has the form '<lhs>-><rhs>' where:

  • <lhs> is a comma-separated list of subscript strings, one per input tensor.
  • <rhs> is the output subscript.
  • The -> separator is required; implicit output subscripts are not supported.
  • Each subscript character must be a single ASCII letter.
  • A letter may appear at most twice within a single input subscript (forming one trace pair). Three or more occurrences raise ValueError.

Contraction order

No contraction-order optimisation is performed. For multi-tensor equations the tensors are contracted strictly from left to right. This may be suboptimal for networks where a different pairing would reduce intermediate tensor sizes. If performance matters, determine the optimal contraction order manually (e.g. by estimating intermediate tensor sizes or using a dedicated contraction-order tool), then rearrange the tensors and equation string accordingly.

Unsupported: Hadamard (batch) indices

A letter that appears in two or more input subscripts and in the output subscript is a Hadamard (batch) index. In NumPy/PyTorch einsum this means element-wise multiplication along that axis with no summation. contract has no such mode, so equations of this kind are not supported (e.g. 'ij,ij->ij').

See Also