Speeding up summation inside comprehension

General Usage Performance
1

I currently do something similar to the code below (though I have axis arrays instead of regular arrays):

A = 1:100
B = 1:200
C = 1:300

x = rand(length(A),length(B),length(C))
y = rand(length(B),length(C))

z = [
    sum(y[b,c]*x[a,b,c] for b=B)
    for a=A, c=C
]

I can see this potentially taking a very long time if my dimensions (A, B and C) get bigger, which they will eventually.

I know that in 2 dimensions I could do a vector/matrix operation which would be faster if I have a multi threaded processor (if I understand correctly), like the following:

x = rand(length(B), length(A))
y = rand(length(B), 1)
z = y' * x

I was wondering if there’s an obvious way to do that for the first example?

2

This is like 4x faster on my machine, could definitely be improved still.

hcat([view(x,:,:,c)*view(y,:,c) for c in C]...)
2 Likes
3

Cool! I didn’t know about the view function. I’ll try and adapt it to my code.

4

This is a tensor contraction, so maybe you can use https://github.com/Jutho/TensorOperations.jl

1 Like
5

Thanks for the recomendation! I implemented tbeason’s use of view and actually got a ~ 20 times speed up, so that’s good enough for me. Makes me want to go through my code and figure out what else I can improve…

1 Like
6

This one isn’t handled by TensorOperations:

TensorOperations.@tensor z[a,c] := y[b,c] * x[a,b,c] # IndexError{String}("non-matching indices

but it can be done with various other packages:

Einsum.@einsum z1[a,c] := y[b,c] * x[a,b,c]
OMEinsum.@ein z2[a,c] := y[b,c] * x[a,b,c]
TensorCast.@reduce z3[a,c] := sum(b) y[b,c] * x[a,b,c]
Tullio.@tullio z4[a,c] := y[b,c] * x[a,b,c]

z5 = reshape(NNlib.batched_mul(x, reshape(y,200,1,300)), 100,300)
z ≈ z1 ≈ z2 ≈ z3 ≈ z4 ≈ z5
7 Likes