Wasserstein implementation does not seem to be fully "batched"

[netw0rkf10w]

Hi @t-vi,

Thanks for sharing your code!

I would like to ask a question regarding your implementation of the Sinkhorn algorithm. You stated that one of the main motivations was to obtain efficient batched computation. However, looking at the code I observe that it only supports the case where the cost matrix is the same across the batch:

def forward(ctx, mu, nu, dist, lam=1e-3, N=100):
        assert mu.dim() == 2 and nu.dim() == 2 and dist.dim() == 2
        bs = mu.size(0)
        d1, d2 = dist.size()
        assert nu.size(0) == bs and mu.size(1) == d1 and nu.size(1) == d2

That is, the shape dist is d1 x d2 instead of bs x d1 x d2. Is this expected?

Thank you in advance for your reply.

[netw0rkf10w]

I also observe that the backward pass does not compute the gradient for dist:

@staticmethod
    def backward(ctx, grad_out):
        return grad_out[:, None] * ctx.log_u * ctx.lam, grad_out[:, None] * ctx.log_v * ctx.lam, None, None, None

which unfortunately prevents learning dist...