Basic Usage

In the following example, we demonstrate the einsum notation for basic tensor operations.

Einsum notation

To specify the operation, the user can either use the @ein_str-string literal or the EinCode object. For example, both the following code snippets define the matrix multiplication operation:

julia> using OMEinsum
julia> code1 = ein"ij,jk -> ik"  # the string literalij, jk -> ik
julia> ixs = [[1, 2], [2, 3]]  # the input indices2-element Vector{Vector{Int64}}:
 [1, 2]
 [2, 3]
julia> iy = [1, 3]  # the output indices2-element Vector{Int64}:
 1
 3
julia> code2 = EinCode(ixs, iy)  # the EinCode object (equivalent to the string literal)1∘2, 2∘3 -> 1∘3

The @ein_str macro can be used to define the einsum notation directly in the function call.

julia> A, B = randn(2, 3), randn(3, 4);
julia> code1(A, B)  # matrix multiplication2×4 Matrix{Float64}:
 -1.36423  -0.422359  0.974233  0.0993386
 -3.82861  -0.266157  2.33438   2.50655
julia> size_dict = OMEinsum.get_size_dict(getixsv(code1), (A, B))  # get the size of the labelsDict{Char, Int64} with 3 entries:
  'j' => 3
  'i' => 2
  'k' => 4
julia> einsum(code1, (A, B), size_dict)  # lower-level function2×4 Matrix{Float64}:
 -1.36423  -0.422359  0.974233  0.0993386
 -3.82861  -0.266157  2.33438   2.50655
julia> einsum!(code1, (A, B), zeros(2, 4), true, false, size_dict)  # the in-place operation2×4 Matrix{Float64}:
 -1.36423  -0.422359  0.974233  0.0993386
 -3.82861  -0.266157  2.33438   2.50655
julia> @ein C[i,k] := A[i,j] * B[j,k]  # all-in-one macro2×4 Matrix{Float64}:
 -1.36423  -0.422359  0.974233  0.0993386
 -3.82861  -0.266157  2.33438   2.50655

Here, we show that the @ein macro combines the einsum notation defintion and the operation in a single line, which is more convenient for simple operations. Separating the einsum notation and the operation (the first approach) can be useful for reusing the einsum notation for multiple input tensors. Lower level functions, einsum and einsum!, can be used for more control over the operation.

For more than two input tensors, the @ein_str macro does not optimize the contraction order. In such cases, the user can use the @optein_str string literal to optimize the contraction order or specify the contraction order manually.

julia> tensors = [randn(100, 100) for _ in 1:4];
julia> optein"ij,jk,kl,lm->im"(tensors...)  # optimized contraction (without knowing the size)100×100 Matrix{Float64}:
  -353.311     118.607     552.614   …    -77.5166    848.71     -1992.82
   438.891     195.403     -58.3405     -1106.11      283.84       236.25
  -439.481     855.166      25.6867      -215.263      -6.50439   1124.3
   781.86     1177.1       838.734        219.782    -166.891     1729.4
   485.832     326.507    -116.64       -1445.52     -506.215     -691.972
  -177.267     741.623   -1496.97    …    273.277   -1729.53      -420.304
   123.895     270.661    -305.19         -28.7883    497.664      128.374
   277.342    -808.237      18.3872      1188.77     -716.617     -506.608
  -439.159    -552.991    -409.185       -698.217   -2276.76      1151.27
    13.2622    331.335    -524.768       -367.778    1266.0       1063.08
     ⋮                               ⋱
  -822.092   -1901.6       -84.1487      -153.312     206.226      500.296
 -1614.83     -809.689    1121.49        1109.03     -201.851    -1810.21
  -122.297    1505.43     -265.441      -1436.04     -369.228    -1210.0
   593.376     271.38     1097.68         651.164      97.4506    -497.513
  1183.73       19.6626   1464.02    …  -1149.64     2029.34      -276.733
   794.647     156.358     316.955        -94.2546  -1377.31       869.838
  -275.497    -923.445     451.18        1369.14     -483.263       89.3984
  -274.825     800.475     131.944        146.44    -1573.74       367.507
  -953.882    -184.085    -873.027        -13.1502    -21.873     2191.79
julia> ein"(ij,jk),(kl,lm)->im"(tensors...)  # manually specified contraction100×100 Matrix{Float64}:
  -353.311     118.607     552.614   …    -77.5166    848.71     -1992.82
   438.891     195.403     -58.3405     -1106.11      283.84       236.25
  -439.481     855.166      25.6867      -215.263      -6.50439   1124.3
   781.86     1177.1       838.734        219.782    -166.891     1729.4
   485.832     326.507    -116.64       -1445.52     -506.215     -691.972
  -177.267     741.623   -1496.97    …    273.277   -1729.53      -420.304
   123.895     270.661    -305.19         -28.7883    497.664      128.374
   277.342    -808.237      18.3872      1188.77     -716.617     -506.608
  -439.159    -552.991    -409.185       -698.217   -2276.76      1151.27
    13.2622    331.335    -524.768       -367.778    1266.0       1063.08
     ⋮                               ⋱
  -822.092   -1901.6       -84.1487      -153.312     206.226      500.296
 -1614.83     -809.689    1121.49        1109.03     -201.851    -1810.21
  -122.297    1505.43     -265.441      -1436.04     -369.228    -1210.0
   593.376     271.38     1097.68         651.164      97.4506    -497.513
  1183.73       19.6626   1464.02    …  -1149.64     2029.34      -276.733
   794.647     156.358     316.955        -94.2546  -1377.31       869.838
  -275.497    -923.445     451.18        1369.14     -483.263       89.3984
  -274.825     800.475     131.944        146.44    -1573.74       367.507
  -953.882    -184.085    -873.027        -13.1502    -21.873     2191.79

Sometimes, manually optimizing the contraction order can be beneficial. Please check Contraction order optimization for more details.

Einsum examples

We first define the tensors and then demonstrate the einsum notation for various tensor operations.

julia> using OMEinsum
julia> s = fill(1)  # scalar0-dimensional Array{Int64, 0}:
1
julia> w, v = [1, 2], [4, 5];  # vectors
julia> A, B = [1 2; 3 4], [5 6; 7 8]; # matrices
julia> T1, T2 = reshape(1:8, 2, 2, 2), reshape(9:16, 2, 2, 2); # 3D tensor

Unary examples

julia> ein"i->"(w)  # sum of the elements of a vector.0-dimensional Array{Int64, 0}:
3
julia> ein"ij->i"(A)  # sum of the rows of a matrix.2-element Vector{Int64}:
 3
 7
julia> ein"ii->"(A)  # sum of the diagonal elements of a matrix, i.e., the trace.0-dimensional Array{Int64, 0}:
5
julia> ein"ij->"(A)  # sum of the elements of a matrix.0-dimensional Array{Int64, 0}:
10
julia> ein"i->ii"(w)  # create a diagonal matrix.2×2 Matrix{Int64}:
 1  0
 0  2
julia> ein"i->ij"(w; size_info=Dict('j'=>2))  # repeat a vector to form a matrix.2×2 Matrix{Int64}:
 1  1
 2  2
julia> ein"ijk->ikj"(T1)  # permute the dimensions of a tensor.2×2×2 Array{Int64, 3}:
[:, :, 1] =
 1  5
 2  6

[:, :, 2] =
 3  7
 4  8

Binary examples

julia> ein"ij, jk -> ik"(A, B)  # matrix multiplication.2×2 Matrix{Int64}:
 19  22
 43  50
julia> ein"ijb,jkb->ikb"(T1, T2)  # batch matrix multiplication.2×2×2 Array{Int64, 3}:
[:, :, 1] =
 39  47
 58  70

[:, :, 2] =
 163  187
 190  218
julia> ein"ij,ij->ij"(A, B)  # element-wise multiplication.2×2 Matrix{Int64}:
  5  12
 21  32
julia> ein"ij,ij->"(A, B)  # sum of the element-wise multiplication.0-dimensional Array{Int64, 0}:
70
julia> ein"ij,->ij"(A, s)  # element-wise multiplication by a scalar.2×2 Matrix{Int64}:
 1  2
 3  4

Nary examples

julia> optein"ai,aj,ak->ijk"(A, A, B)  # star contraction.2×2×2 Array{Int64, 3}:
[:, :, 1] =
 68   94
 94  132

[:, :, 2] =
  78  108
 108  152
julia> optein"ia,ajb,bkc,cld,dm->ijklm"(A, T1, T2, T1, A)  # tensor train contraction.2×2×2×2×2 Array{Int64, 5}:
[:, :, 1, 1, 1] =
  9500  14564
 21604  33420

[:, :, 2, 1, 1] =
 11084  17012
 25204  39036

[:, :, 1, 2, 1] =
 13644  20916
 31028  47996

[:, :, 2, 2, 1] =
 15932  24452
 36228  56108

[:, :, 1, 1, 2] =
 13214  20258
 30050  46486

[:, :, 2, 1, 2] =
 15414  23658
 35050  54286

[:, :, 1, 2, 2] =
 19430  29786
 44186  68350

[:, :, 2, 2, 2] =
 22686  34818
 51586  79894

Computation Backends

OMEinsum supports multiple backends for tensor contractions. The backend determines how the underlying computation is performed.

Available Backends

Backend	Description	Best For
`DefaultBackend()`	BLAS/CUBLAS via reshape/permute	General use, matrix operations
`CuTensorBackend()`	NVIDIA cuTENSOR	GPU tensor network contractions

Changing Backends

julia> get_einsum_backend()  # check current backendDefaultBackend()
julia> set_einsum_backend!(DefaultBackend())  # set to defaultDefaultBackend()

For GPU acceleration with cuTENSOR, see CUDA Acceleration.