PyTorch doubly stochastic normalisation of 3D tensor
I'm trying to implement double stochastic normalisation of an N x N x P tensor as described in Section 3.2 in Gong, CVPR 2019. This can be done easily in the N x N case using matrix operations but I am stuck with the 3D tensor case. What I have so far is
def doubly_stochastic_normalise(E):
    &quot;&quot;&quot;E: n x n x f&quot;&quot;&quot;
    E = E / torch.sum(E, dim=1, keepdim=True)  # normalised across rows
    F = E / torch.sum(E, dim=0, keepdim=True)  # normalised across cols
    E = torch.einsum('ijp,kjp-&gt;ikp', E, F)
return E

but I'm wondering if there is a method without einsum.

In this setting, you can always fall back to using torch.matmul (batched matrix multiplication to be more precise). However, this requires you to transpose the axis. Recall the matrix multiplication for two 3D inputs, in einsum notation, it gives us:
bik,bkj-&gt;bij

Notice how the k dimension gets reduces. To get to this setting, we need to transpose the inputs of the operator. In your case we have:
ijp ? kjp -&gt; ikp
 ↓     ↓      ↑
pij @ pjk -&gt; pik

This translates to:
&gt;&gt;&gt; (E.permute(2,0,1) @ F.permute(2,1,0)).permute(1,2,0)
#       ijp ➝ pij          kjp ➝ pjk      pik ➝ ikp

You can argue your method is not only shorter but also a lot more readable. I would therefore stick with torch.einsum. The reason why the einsum operator is so useful here is because you can perform axes transpositions on the fly.
