[FR] add Gumbel distribution

[fehiepsi]

As requested @tbsexton at #545, it would be nice to have Gumbel distribution in NumPyro.

References: wikipedia, pytorch dist, jax sampler.

[rtbs-dev]

Also commenting to add info on gumbel-softmax, which kicked off this discussion (continuous approximation of categorical distribution).

Paper, (deprecated?) Pytorch version with straight-through forward-mode support

[rtbs-dev]

Wow, thanks to @daydreamt! That was fast. I'd be happy to close this, or re-name it as "Gumbel-Softmax" if that functionality is also wanted.

Based on the pytorch link above, it seems like they are making use of some pytorch-specific detachments, etc. to implement the straight-through version (argmax in the fwd/sampling direction, softmax in the backward/gradient direction, if I'm understanding correctly). Not sure how such a thing could be done with the JAX backend.

[fehiepsi]

@tbsexton A similar version can be implemented using jax.lax.stop_gradient, which plays the role of PyTorch detach. However, I am not sure how useful it is to solve your problem. Could you open a new issue which sketches your ideas and your feature request? FYI, we have RelaxedOneHotCategoricalStraightThrough distribution in Pyro, which can be ported to NumPyro without much difficulty.

I would like to close this issue for now and discuss on the new one. :)