I am working on a small package to implement the Sinkhorn algorithm for optimal transport.
Since this is frequently used as a loss function for generative models, my goal is to have an implementation that is generic enough to work on GPU or CPU, and to work with normal arrays as well as autodiff code (currently Flux TrackedArray).
All of the relevant code is here; it’s a pretty simple algorithm to implement. I am wondering, though, if there’s anything obviously dubious or non-idiomatic in the way I’ve written it.
(Some of the allocations may appear unnecessary, but in fact to work with TrackedArrays I am pretty sure it does have to be written this way, rather than mutating an array and avoiding the allocations.)