Numpy dot product along specific axes

user_4830338 0
Viewed 429

I have a 512x512 image array and I want to perform operations on 8x8 blocks. At the moment I have something like this:

output = np.zeros(512, 512)

for i in range(0, 512, 8):
    for j in rangerange(0, 512, 8):
        a = input[i:i+8, j:j+8]
        b = some_other_array[i:i+8, j:j+8]
        output[i:i+8, j:j+8] = np.dot(a, b)

where a & b are 8x8 blocks derived from the original array. I would like to speed up this code by using vectorised operations. I have reshaped my inputs like this:

input = input.reshape(64, 8, 64, 8)
some_other_array = some_other_array.reshape(64, 8, 64, 8)

How could I perform a dot product on only axes 1 & 3 to output an array of shape (64, 8, 64, 8)?

I have tried np.tensordot(input, some_other_array, axes=([0, 1], [2, 3])) which gives the correct output shape, but the values do not match the output from the loop above. I've also looked at np.einsum but I haven't come across a simple example with what I'm trying to achieve.

numpy python
2 Answers

As you suspected, np.einsum can take care of this. If input and some_other_array have shapes (64, 8, 64, 8), then if you write

output = np.einsum('ijkl,ilkm->ijkm', input, some_other_array)

then output will also have shape (64, 8, 64, 8), where matrix multiplication (i.e. np.dot) has been done only on axes 1 and 3.

The string argument to np.einsum looks complicated, but really it's a combination of two things. First, matrix multiplication is given by jl,lm->jm (see e.g. this answer on einsum). Second, we don't want to do anything to axis 0 and 2, so for them I just write ik,ik->ik. Combining the two gives ijkl,ilkm->ijkm.

user_11030440

They'll work if you reorder them a bit. If input and some_other_array are both shaped (64,8,64,8), then:

input = input.transpose(0,2,1,3)
some_other_array = some_other_array.transpose(0,2,1,3)

This will reorder them to 64,64,8,8. At this point you can compute a matrix multiplication. Do note that you need matmul to compute the block products, and not dot, which will try to multiply the entire matrices.

output = input @ some_other_array
output = output.transpose(0,2,1,3)
output = output.reshape(512,512)
user_102297540
Related