"non-matching indices" error using TensorOperations

General Usage
1

I’m trying to simplify some code that does a lot of dropdims and mapslices, so I’m looking at some packages that do Einstein notation. I tried out Einsum.jl which seems pretty intuitive, but I also wanted to try out TensorOperations.jl, which seems to have higher performance and some bells and whistles like GPU support.

So I figured I’d try out TensorOperations but I’m getting an error that I don’t really understand.

Here’s a MWE for the operation I’m trying to do:

a = randn(3,4,5)
b = randn(5,3)
@tensor c[k, n] := a[k, n, c] * b[c,k]

I expected this to translate to something like:

c = [sum(a[k, n, :] .* b[:, k]) for k in 1:size(a,1), n in 1:size(a,2)]

Which is what @einsum does. But this gives the error:

IndexError{String}("non-matching indices between left and right hand side: \$(Expr(:(:=), :(var\"##835\"[k, n]), :(var\"##833\"[k, n, c] * var\"##834\"[c, k])))")

Can anyone help me understand why this is valid for @einsum but not @tensor? My mental model (based on the Einsum docs) is that variables introduced in the right-hand-side of an expression are summed over, but TensorOperations seems to have more restrictions on what’s a valid expression.

2

The @einsum macro writes loops as you suggest, which isn’t fussy about what’s on the right (including scalar functions). But @tensor decomposes into a sequence of operations, and doesn’t consider this one to be valid: it’s some kind of batched matrix multiplication. However @ein has that as one of its basic operations, so you can do this:

using OMEinsum
@ein c[k, n] := a[k, n, c] * b[c,k]

(It also has a fallback implementation more like @einsum, but this is not being used here, you can check with allow_loops(false).)

3

Thanks for the explanation. From what I can tell it seems like OMEinsum looks really promising and has a lot of the nice features of both Einsum and TensorOperations.

Unfortunately it doesn’t seem to support conj in the expressions, which makes it tricky to represent complex-valued operations. I filed an issue, and for now will probably stick with Einsum unless the performance actually becomes a bottleneck in practice. Most of the time I just want the nice syntax for munging multi-dimensional data and performance isn’t super critical.

4

With TensorOperations you can only saturate indices completely, so in your example index k cannot appear on the left-hand side.
One workaround this to contract with a copy tensor

δ = zeros(3,3,3);
for i in first(axes(δ))
    δ[i,i,i] = 1
end
@tensor c[l, n] := a[j, n, c] * b[c, k] * δ[j, k, l]
1 Like
5

For what it’s worth, I think @mcabbott is being a bit modest by not mentioning his package TensorCast.jl which, while not always as performant as TensorOperations.jl, has a lot of very cool bells and whistles as well as being way more flexible than Einsum or TensorOperations as far as I’m concerned:

julia> using TensorCast

julia> a = randn(ComplexF64, 3, 4, 5); b = randn(ComplexF64, 5, 3);

julia> @reduce c[k, n] := sum(l) a[k, n, l] * b[l, k]
3×4 Array{Complex{Float64},2}:
 0.828145+2.201im     -0.381748+1.32616im      0.62934-0.695514im  -0.468997+1.43541im
 -3.70024-3.0064im     -2.36539-2.51518im      1.26878-0.992282im   0.199802-1.22858im
 -0.78809-0.826066im  -0.969338-0.0449697im  0.0591432-1.92419im   -0.396283-0.943951im

julia> @reduce c[k, n] := sum(l) a[k, n, l] * conj(b)[l, k]
3×4 Array{Complex{Float64},2}:
  1.48318+0.0390588im  -0.071153-0.303195im  -0.259096+1.22923im   0.74586-1.00576im
 -2.77538+0.383936im     1.15156+3.0815im      0.58602-2.90407im  -0.81003+2.48695im
  0.71891+3.54748im    -0.690237+0.824959im  -0.756242+1.03495im   1.25198-0.632437im

now let’s try on the GPU:

julia> using CuArrays; cua = cu(a); cub = cu(b);

julia> @reduce cuc[k, n] := sum(l) cua[k, n, l] * conj(cub)[l, k]
3×4 CuArray{Complex{Float64},2,CuArray{Complex{Float64},3,Nothing}}:
  1.48318+0.0390588im  -0.071153-0.303195im  -0.259096+1.22923im   0.74586-1.00576im
 -2.77538+0.383936im     1.15156+3.0815im      0.58602-2.90407im  -0.81003+2.48695im
  0.71891+3.54748im    -0.690237+0.824959im  -0.756242+1.03495im   1.25198-0.632437im
5 Likes
6

Thanks Mason! I will add however that you can move the conjugation outside, so that it gets done in one broadcast:

julia> @pretty @reduce c[k, n] := sum(l) a[k, n, l] * b[l, k]'
begin
    local turkey = orient(PermuteDims(b), (:, *, :))
    c = dropdims(sum(@__dot__(a * adjoint(turkey)), dims=3), dims=3)
end

Somewhat surprisingly, this approach of broadcasting a huge array and summing it is faster than the dedicated batched matmul (plus permutedims) here, even for quite large arrays. Although both are beaten by naiive loops:

julia> f_ome(a,b) = @ein c[k, n] := a[k, n, c] * conj(b)[c,k];
julia> f_cast(a,b) = @reduce c[k, n] := sum(l) a[k, n, l] * b[l, k]';
julia> f_ein(a,b) = @einsum c[k, n] := a[k, n, c] * conj(b[c,k]);

julia> a = randn(ComplexF64, 300, 400, 500); b = randn(ComplexF64, 500, 300);

julia> f_ome(a,b) ≈ f_cast(a,b) ≈ f_ein(a,b)
true

julia> @btime f_ome($a, $b);
  1.845 s (243 allocations: 921.49 MiB)

julia> @btime f_cast($a, $b);
  649.614 ms (26 allocations: 919.65 MiB)

julia> @btime f_ein($a, $b);
  261.044 ms (2 allocations: 1.83 MiB)
2 Likes
7

Awesome, I’m playing around with TensorCast and it’s indeed very slick.

Is there a way to add or subtract a constant outside the reduction for each element in the map? Currently I’m using the approach described in the recursion section where I use both a @cast and a reduce but it seems like maybe I’m making things over-complicated.

Here’s what I have now:

@cast w[k] := (@reduce [k] := sum(c) abs(vs[k, c])) - 1
8

I see, you want something like w_k = -1 + Σ_c |vs_k,c|. No, there is not I’m afraid, except doing it in two steps. It will always sum the entire expression appearing after sum(...).