Tensor Contractions for N-dimensional tensors

General Usage Performance
, ,
1

I am working with exponentially sized tensors (exponential in a given variable N, for my case to the basis 4). Up to now, I stored the tensors linearly indexed, that is in a vector from 1:4^N. Contractions where therefore straightforward using packages like TensorOperations.jl for example.
In a next step I want to utilize the sparsity of the matrices I wish to contract the vector with. It is therefore beneficial (on paper) to write the vector as an N dimensional tensor like this: [1:4, 1:4, …, 1:4], because then I only need contraction with ONE of the N 1:4 blocks (the others drop out due to the sparsity).
There are a few ways to achieve this. I could store the vector linearly indexed and then “by hand” (i.e. using the digits function with base 4) find the relevant indices that I need to contract. This would be my way to go, but here it is said that this might not be a good idea.
The other way would be to store the vector like this [1:4, 1:4, …, 1:4] and then use TensorOperations.jl base functions contracttensors!(…) to say which indices to contract. This feels kind of hacky because you need an N item string and in that string you need to set the k-th element such that it contracts as desired. I am worried about performance, and this surely does not sound optimal.
However, I could not find a better option. Ideally there’d exist a function contract(A, B, 2, 5) to contract tensors A and B along dimensions 2 and 5 respectively. Did anyone ever run into a similar problem and if so, how did you solve it? Thanks a lot!

2

Using TensorOperations.tensorcontract sounds reasonable. For concreteness, I believe a function like this one is roughly what you meant (correct me if I am wrong).

function contract(A,B,kA,kB)
    IA = collect(-ndims(A):-1)
    IB = collect(1:ndims(B))
    IA[kA] = 0
    IB[kB] = 0
    return tensorcontract(A,IA,B,IB)
end

I agree that creating IA and IB seems wasteful, but actually it’s not that bad. The actual contraction scales exponentially in N while creating IA and IB scales only linearly in N, so most likely the cost of creating IA and IB will be negligible compared to the contraction.

1 Like
3

Yes, I came up with the same thing but without the nice form for IA and IB. Thanks a lot, I’ll be using this!

4

This might be something for @mcabbott’s Tullio.jl ?

5

Here’s a quick test with N=4; TensorOperations will do better at larger N, and Einsum better at smaller N:

julia> A = rand(4,4,4,4); B = rand(4,4,4,4);

julia> using TensorOperations, Einsum, Tullio, LoopVectorization
julia> ten33(A,B) = @tensor C[a,b,c,x,y,z] := A[a,b,s,c] * B[x,y,s,z]
julia> ein33(A,B) = @einsum C[a,b,c,x,y,z] := A[a,b,s,c] * B[x,y,s,z]
julia> tul33(A,B) = @tullio C[a,b,c,x,y,z] := A[a,b,s,c] * B[x,y,s,z]  tensor=false

julia> @btime ten33($A,$B);
  12.612 μs (80 allocations: 36.94 KiB)

julia> @btime ein33($A,$B);
  11.752 μs (3 allocations: 32.20 KiB)

julia> @btime tul33($A,$B);
  3.596 μs (14 allocations: 32.56 KiB)

To make a function which accepts the dimensions, for Einsum you can make a generated function contract(A, B, Val(3), Val(3)) which writes the macro. But Tullio defines functions and thus doesn’t play well with @generated, unfortunately.

You could write a function which always contracts the 2nd index, and then reshapes. Something like this:

sizeA3 = prod(size(A,d) for d in 1:kA-1), size(A,kA), prod(size(A,d) for d in kA+1:ndims(A))
A3 = reshape(A, sizeA3)
B3 = reshape(B, ...) # similar

@tullio C4[a,b,x,y] := A3[a,s,b] * B3[x,s,y] avx=true tensor=false

sizeCA = filter(!iszero, ntuple(d->d==kA ? 0 : size(A,d), ndims(A)))
C = reshape(C4, tuple(sizeCA..., sizeCB...))

Edit: a bit surprisingly, the function above is much quicker than the one with @tensor macro.

julia> @btime contract($A,$B,3,3);
  5.962 μs (46 allocations: 39.52 KiB)

julia> function contract2(A,B,kA,kB)
           IA = ntuple(d -> d==kA ? 0 : -d, ndims(A))
           IB = ntuple(d -> d==kB ? 0 : +d, ndims(B))
           return TensorOperations.tensorcontract(A,IA,B,IB)
       end

julia> @btime contract2($A,$B,3,3);
  5.366 μs (47 allocations: 39.47 KiB)
2 Likes
6

That’s the effect of the cache for temporaries being used with the @tensor macro. It is helpful for large tensor contractions involving several temporaries, but for small test cases like this, it is better to run with TensorOperations.disable_cache().

1 Like