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 literal
ij, jk -> ik
julia> ixs = [[1, 2], [2, 3]] # the input indices
2-element Vector{Vector{Int64}}: [1, 2] [2, 3]
julia> iy = [1, 3] # the output indices
2-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 multiplication
2×4 Matrix{Float64}: -2.61076 -1.95243 -0.656993 1.43942 -2.17537 -2.84521 -2.38323 3.31437
julia> size_dict = OMEinsum.get_size_dict(getixsv(code1), (A, B)) # get the size of the labels
Dict{Char, Int64} with 3 entries: 'j' => 3 'i' => 2 'k' => 4
julia> einsum(code1, (A, B), size_dict) # lower-level function
2×4 Matrix{Float64}: -2.61076 -1.95243 -0.656993 1.43942 -2.17537 -2.84521 -2.38323 3.31437
julia> einsum!(code1, (A, B), zeros(2, 4), true, false, size_dict) # the in-place operation
2×4 Matrix{Float64}: -2.61076 -1.95243 -0.656993 1.43942 -2.17537 -2.84521 -2.38323 3.31437
julia> @ein C[i,k] := A[i,j] * B[j,k] # all-in-one macro
2×4 Matrix{Float64}: -2.61076 -1.95243 -0.656993 1.43942 -2.17537 -2.84521 -2.38323 3.31437
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}: -1996.2 -688.844 52.191 … 1717.8 -759.934 -364.637 1147.1 -270.668 -644.16 192.212 -977.341 -15.4255 -1038.38 -991.216 -355.294 851.935 -435.09 1458.44 -440.975 -1195.49 -342.819 -87.772 239.591 512.456 -349.853 129.605 25.1695 -911.206 -588.266 -824.554 1571.02 697.185 -321.533 … 1265.89 503.258 522.084 -933.436 -282.084 -865.232 -459.03 -866.671 -292.52 367.692 220.019 -1085.33 -1276.33 -652.851 -907.663 -287.869 1343.33 -1000.09 -652.037 -234.802 -129.831 332.608 771.604 -636.148 -78.0908 457.035 14.4169 ⋮ ⋱ 1267.3 221.675 -1137.79 -296.931 -1293.52 1609.43 -642.229 -951.205 -1281.86 204.266 -1250.92 150.965 657.418 -313.463 1409.44 -936.032 -133.6 -126.27 340.744 -419.117 792.555 832.515 695.239 2136.14 1393.96 1037.57 -1105.43 … -2375.31 -725.771 -2029.0 801.296 -460.749 592.687 -879.581 -116.868 -835.627 1413.02 1404.45 527.074 -1490.64 -162.535 969.293 -368.059 1098.93 2335.78 -326.061 355.318 346.483 426.73 -122.229 -498.937 -1660.01 -1515.6 -1019.39
julia> ein"(ij,jk),(kl,lm)->im"(tensors...) # manually specified contraction
100×100 Matrix{Float64}: -1996.2 -688.844 52.191 … 1717.8 -759.934 -364.637 1147.1 -270.668 -644.16 192.212 -977.341 -15.4255 -1038.38 -991.216 -355.294 851.935 -435.09 1458.44 -440.975 -1195.49 -342.819 -87.772 239.591 512.456 -349.853 129.605 25.1695 -911.206 -588.266 -824.554 1571.02 697.185 -321.533 … 1265.89 503.258 522.084 -933.436 -282.084 -865.232 -459.03 -866.671 -292.52 367.692 220.019 -1085.33 -1276.33 -652.851 -907.663 -287.869 1343.33 -1000.09 -652.037 -234.802 -129.831 332.608 771.604 -636.148 -78.0908 457.035 14.4169 ⋮ ⋱ 1267.3 221.675 -1137.79 -296.931 -1293.52 1609.43 -642.229 -951.205 -1281.86 204.266 -1250.92 150.965 657.418 -313.463 1409.44 -936.032 -133.6 -126.27 340.744 -419.117 792.555 832.515 695.239 2136.14 1393.96 1037.57 -1105.43 … -2375.31 -725.771 -2029.0 801.296 -460.749 592.687 -879.581 -116.868 -835.627 1413.02 1404.45 527.074 -1490.64 -162.535 969.293 -368.059 1098.93 2335.78 -326.061 355.318 346.483 426.73 -122.229 -498.937 -1660.01 -1515.6 -1019.39
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) # scalar
0-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