Inaccurate gradient estimation for Gumbel-Sinkhorn

[adamoyoung]

I implemented a Gumbel-Sinkhorn layer as a CVXPY layer, and compared it with a baseline implementation using the Sinkhorn algorithm. While the forward passes of both layers seem similar, the backward passes differ. Specifically, it seems like the CVXPY layer is not able to provide accurate gradient esimates. I am looking for advice as to why it is not working properly.

I have detailed my experiments in this Google colab gist. Any feedback would be appreciated. Thanks!

[AxelBreuer]

Hi @adamoyoung ,
I am not familiar with Gumbel-Sinkhorn layers, but I quickly went through your Google colab and realized that this layer uses cp.log_sum_exp(x) and cp.exp(x).
I have recently fixed a bug in exponential cones (cvxgrp/diffcp#59); hence it could be woth to update to the lastest diffcp and relaunch you tests.
With a bit of luck, your code will work now.