Is a symmetric softmax output for 2d array possible?

vision
1

Hi,
I have a (n,1,k,k) tensor; n is the batch size and k is number of different objects in a picture.
I’m trying to create a network that produces the similarity between different objects in a given picture.
I was wondering if there’s any method that produces a symmetric softmax. (maybe something that produces a 'diagonal-like matrix" where each object has the most similarity with itself)

This is what my forward method looks like, but the output isn’t symmetrical.
I also tried using SoftMax2D but all elements became 1.

def forward(self, v_pc):
        vdot_pc = self.pg_layers(v_pc)
        vdot_pc=vdot_pc.squeeze() # (:,100,100) (k*k)

        vdot_pc = torch.add(vdot_pc, vdot_pc.transpose(1, 2)) 

        out = nn.functional.softmax(vdot_pc, dim = 2)

        return out

I could try making the output symmetrical by adding it to its transpose, but I’m not sure which dimension would be appropriate for this task.

I’d appreciate your help.