Non-Local Affinity Row

Problem

Compute one row of the non-local block affinity matrix.

Flatten the H x W x C feature map into positions, take the dot product of the query position with every position, apply softmax over the scores and return the flattened list rounded to 4 decimal places.

Test cases

case 1

in: [[[[1,0],[0,1]],[[1,1],[1,1]]],0]

expected: [0.2969,0.1092,0.2969,0.2969]

case 2

in: [[[[1,0],[0,1]],[[1,1],[1,1]]],3]

expected: [0.1345,0.1345,0.3655,0.3655]

case 3

in: [[[[0,0],[0,0]],[[0,0],[0,0]]],1]

expected: [0.25,0.25,0.25,0.25]

Starter codePython

def nonlocal_affinity(feature, query_index):
    # Your code here
    pass

3 test cases · run your code against them in the editor