Speeding up summation inside comprehension

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