Einstein convention for slicing or how to view diagonals in tensors

If you like np.einsum, then you can do these things as follows:

julia> using TensorCast

a = zeros(2,3,1);

julia> @cast b[k,j,i] := a[i,j,k]
1×3×2 PermutedDimsArray(::Array{Float64,3}, (3, 2, 1)) with eltype Float64:
[:, :, 1] =
 0.0  0.0  0.0

[:, :, 2] =
 0.0  0.0  0.0

julia> b[1,1,:] = 1:2
1:2

julia> a
2×3×1 Array{Float64,3}:
[:, :, 1] =
 1.0  0.0  0.0
 2.0  0.0  0.0

julia> c = zeros(2,2);

julia> @cast d[i] := c[i,i]
2-element view(::Array{Float64,1}, 1:3:4) with eltype Float64:
 0.0
 0.0

julia> d .= [7,9];

julia> c
2×2 Array{Float64,2}:
 7.0  0.0
 0.0  9.0

julia> using LinearAlgebra

julia> view(c, diagind(c))
2-element view(::Array{Float64,1}, 1:3:4) with eltype Float64:
 7.0
 9.0

But you can also call PermutedDimsArray and diagind by yourself. Or write loops.

2 Likes