Feature corelation

I have a feature map of size NxCxWxH, where N is the batch size, C is the number of channels, and W and H are width and the height respectively.

Lets consider features (a_i) from this feature map. Each feature has a dimension of Cx1x1. I need to find out the following for each of the feature: \sum_j a_i*a_j.

So the output I am expecting is Nx1xWxH, where each element represents the sum of the dot product between the feature at that location and all the other features. How to do it?

Create a map sized CxWxH of the feature a_i repeated WxH times, then calculate a dot product along the channel dimension between that and your original feature map of size CxWxH (let’s call it A):

A_rep = a_i.repeat((1,W,H))
result = (A_rep*A).sum(axis=0)

I ignored N for simplicity but the concept is similar with N>1