Solved by googlebot in post #2 
              
              
                You can imagine sort() as doing two operations: 
indices = argsort(x,dim) 
y = gather(x,indices,dim) #x[indices] may work too 
argsort() defines non-trainable permutation 
gather() applies this permutation. It is a differentiable operation - backward pass permutes gradients back to match their origi…