Efficiently creating array of dimension 3 from element-wise matrix addition

General Usage
1

I have to perfrom an operation, in which I create an array with dimension 3 from two matrices in the following way:

C[i,j,k] = A[i,j] + B[j,k]

I can do this of course with a loop but this seems by far not the best way to compute this. Is there a more efficient way?

2

Loops are not inefficient, on the contrary any efficient solution is probably equivalent to a loop under the hood!

You could write this directly with an array comprehension:

A = rand(3,4)
B = rand(4,5)

C = [A[i,j] + B[j,k] for i in axes(A,1),
                         j in axes(A,2),
                         k in axes(B,2)]

but this is a perfect job for an Einstein summation macro, for example with Tullio:

using Tullio

@tullio C[i,j,k] := A[i,j] + B[j,k]

Hard to make this more readable, and it should be very fast :slight_smile:

5 Likes
3

FWIW, here is an alternative using TensorCast, which should be more lightweight to load and to run the first time than Tullio:

using TensorCast
@cast C[i,j,k] := A[i,j] + B[j,k]
1 Like