torch.einsum
torch.einsum(equation, *operands) -> Tensor
Sums the product of the elements of the input "operands" along
   dimensions specified using a notation based on the Einstein
   summation convention.
Einsum allows computing many common multi-dimensional linear
   algebraic array operations by representing them in a short-hand
   format based on the Einstein summation convention, given by
   "equation". The details of this format are described below, but the
   general idea is to label every dimension of the input "operands"
   with some subscript and define which subscripts are part of the
   output. The output is then computed by summing the product of the
   elements of the "operands" along the dimensions whose subscripts
   are not part of the output. For example, matrix multiplication can
   be computed using einsum as torch.einsum("ij,jk->ik", A, B).
   Here, j is the summation subscript and i and k the output
   subscripts (see section below for more details on why).