Hi @tom,
nice work !!!
I found a very simple implementation of the sinkhorn-knopp matrix normalisation, we were looking at a while ago. It’s in MATLAB, but it’s very understandable,

  
      mathworks.com
  
  
    

Sinkhorn-Knopp algorithm for matrix normalisation - File Exchange - MATLAB... 7

Normalises a matrix so that the row and column sums are all unity


  
  
    
    
  
  


Hopefully I’ll try to implement it in torch this weekend - that should be a lot of fun, and useful too.
All the best,
Aj