Einstein convention for slicing or how to view diagonals in tensors

I have looked Einsum.jl and TensorOperations.jl to replicate following behavior np.einsum:

Views returned from einsum are now writeable whenever the input array is writeable. For example, np.einsum('ijk...->kji...', a) will now have the same effect as np.swapaxes(a, 0, 2) and np.einsum('ii->i', a) will return a writeable view of the diagonal of a 2D array.

This seems not yet possible also I took a look @views but that seemed cumber some too. Any suggestions? The main questions is how to define views of various digaonal and populate various diagonals in a tensor?

1 Like

Unlike in NumPy, it is perfectly reasonable and performant to simply write a loop in Julia if you want to populate some diagonal of a multidimensional array.

1 Like

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

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:

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:

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


Thanks a lot!

Somehow the above suggestion throws an error in the above setting.

## tensorcasts
using TensorCast

@cast A[i]:= B[i,i]

@cast A[i,l]:= B[i,i,l,$a]
@castA[i,l]:= B[i,i,l,l]
ERROR: LoadError: index i repeated in [i, i, l]
    @cast A[i, l] := B[i, i, l, $(Expr(:$, :a))]  

Any help greatly appreciated

There are lots of little clever tricks that you can use for indexing with Julia without going to Einstein notations.

For example, linear indexing can enable a simple diagonal in a 2-d matrix


Or you can broadcast CartesianIndex to a similar effect:


The above two forms will compose nicely with views. You can also simply use broadcast directly — but you need to “protect” the original matrix from participating:

getindex.(Ref(A), axes(A)...)

While that form won’t work with views, it will fuse with other broadcasted expressions and may be the most efficient of all (depending upon what else you’re doing) — and you can easily transform it into a setindex! call, too.

Einstein notation is great if you’re familiar with it — and there are several packages that expose this notation — but there are a bunch of other ways to do this, too!


In TensorCast, repeated indices only work on matrices right now, essentially because Diagonal exists but higher versions don’t:

julia> v = 1:3

julia> @cast d[i,i] := v[i]
3×3 LinearAlgebra.Diagonal{Int64,UnitRange{Int64}}:
 1  ⋅  ⋅
 ⋅  2  ⋅
 ⋅  ⋅  3

julia> m = rand(2,3);

julia> @cast z[i,i, j,j] := m[i,j]
ERROR: LoadError: index i repeated in [i, i, j, j]

(Docs on this here.) I was going to generalise this (and diagind / 1:size(A,1)+1:end likewise) but haven’t got around to it.

Einsum is quite happy about repetition at least on the right:

julia> using Einsum

julia> B = randn(2,2,2,2);

julia> @einsum A[i,l] := B[i,i,l,l]
2×2 Array{Float64,2}:
 2.0973     -1.11089
 0.0396539  -0.168977

Thanks a lot but Einsum does a summation as opposed to a cast or view which is what I am looking for.

It does make a new array, but doesn’t sum anything here. But you can make a similar view with CartesianIndex tricks:

julia> @macroexpand1( @einsum A[i,l] := B[i,i,l,l] ) |> MacroTools.prettify
    local baboon = promote_type(eltype(B))
    A = Array{baboon}(undef, size(B, 1), size(B, 3))
    @assert size(B, 3) == size(B, 4)
    @assert size(B, 1) == size(B, 2)
    let i, l
        @inbounds for l = 1:size(B, 3)
                for i = 1:size(B, 1)
                    A[i, l] = B[i, i, l, l]

julia> B = randn(2,2,2,2);

julia> @einsum A[i,l] := B[i,i,l,l];

julia> Aview = view(B, CartesianIndex.(axes(B,1), axes(B,2)), CartesianIndex.(axes(B,3), axes(B,4)))
2×2 view(::Array{Float64,4}, CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(2, 2)], CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(2, 2)]) with eltype Float64:
 0.442564  -0.325475
 0.483786   0.0225074

julia> Aview .= 99;

julia> B
2×2×2×2 Array{Float64,4}:
[:, :, 1, 1] =
 99.0       -0.595575
  0.670441  99.0

[:, :, 2, 1] =
 1.23127  0.0496867
 0.4866   0.0557377

[:, :, 1, 2] =
  0.946728  -0.0395328
 -0.36525   -0.675743

[:, :, 2, 2] =
 99.0       0.710038
 -2.40632  99.0
1 Like

Thanks a lot but I don;t know how this generalises to


or more compliated expressions…Any help appreciated

You can probably do this with CartesianIndex.(axes(B,3), transpose(axes(B,4)), axes(B,5)), but these are pretty exotic objects. Maybe there’s a simpler way to do things? I don’t think np.einsum makes a view here.