TODO:

• Flake8
• Add new tests
• Decide on whether to use torch.Distributions-api or .exponential_() and then remove commented code

  pytorch/torch/nn/functional.py

  Line 1043 in e827a7e

  # dist = torch.distributions.RelaxedOneHotCategorical(

Only error with original test was caused by flake8 (hindering other pull requests to merge currently)

$ flake8
./torch/nn/functional.py:1173:121: E501 line too long (130 > 120 characters)
./torch/nn/functional.py:1247:121: E501 line too long (130 > 120 characters)
./torch/nn/functional.py:1248:121: E501 line too long (123 > 120 characters)

I think this is ready to be merged.

• It worked with the original tests,
• the new tests are more rigorous.
• the api has not changed so no breaking changes.

Here is a comparison between implementations with full training example https://gist.github.com/ragulpr/1f88aee9fcc585600280ba2d6f323368

Or nbviewer if above link not working:
https://nbviewer.jupyter.org/urls/gist.githubusercontent.com/ragulpr/1f88aee9fcc585600280ba2d6f323368/raw/ed5ae3e85064768886efd656b50261ed2f590633/gumbel_softmax.ipynb

Anyone who could review? I know @soumith reviewed original PR 👍

It would be nice if we keep implementations of RelaxedOneHotCategorical.rsample and F.gumbel_softmax consistent.

Thanks @alicanb I've been looking into how that would be possible. I think it's two things; consistent in values or consistent in style. Long story short, the line-by-line translation using reference implementation

Alt a)

    uniforms = clamp_probs(torch.rand(logits.shape, dtype=logits.dtype, device=logits.device))
    gumbels = -((-(uniforms.log())).log()) # -log(exponential) = gumbel(0,1)
    scores = (logits + gumbels) / tau # gumbel(logits,tau)
    y_soft = scores - scores.logsumexp(dim=dim, keepdim=True)
    # need to apply transform(x) as it's a TransformedDistribution
    # https://github.com/pytorch/pytorch/blob/99ce499bfecd4ed0e87cf60d87ee483c6383d95c/torch/distributions/transformed_distribution.py#L99
    y_soft = y_soft.exp() 

14.4 s ± 86 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

But this is going to be slow and logsumexp unnecessary. To speed this up and leave a smaller memory footprint by reusing gumbels variable we get below. This is still consistent in values (generates same thing)

Alt b)

    # merge logsumexp and exp into softmax, reuse variable names
    uniforms = clamp_probs(torch.rand(logits.shape, dtype=logits.dtype, device=logits.device))
    gumbels = -((-(uniforms.log())).log()) # -log(~Exp(1)) = ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)

11.5 s ± 130 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Third alternative is to impose the more legible 'style' of relaxed_categorical rsample implementation but use exponential_() and .new(). This is slightly different in values as it doesn't clamp_probs as aggressively.

Alt c)

    gumbels = -logits.new(logits.shape).exponential_().log() # ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)  # Gumbel softmax noise

11.1 s ± 58.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

For reference, this was what I committed;

Alt d)

    y_soft = logits.new(logits.shape)
    y_soft = (logits - y_soft.exponential_().log()) / tau  # Gumbel noise
    y_soft = y_soft.softmax(dim)  # Gumbel softmax noise

10.6 s ± 59 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

And using finally using RelaxedOneHotCategorical without inlining

    dist = torch.distributions.RelaxedOneHotCategorical(
        temperature=tau, logits=logits, validate_args=False)
    y_soft = dist.rsample()

21.2 s ± 326 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In the end, if we want to mimic in values I propose going with alt b) otherwise alt c)

@alicanb Do you have an opinion about the implementation strategies?

I'm not sure what's going on with the failures, because the mentioned code doesn't show up in the PR at all. Pushed a merge to master and trying again.

Unfortunately, changes in JIT mean that this patch no longer works, because clamp_probs is not scriptable.

@wanchaol could you take a look at this?

Yes this will make jit standard lib break , @ragulpr why we ended up choosing Alt b) instead of Alt c) given that the Alt c) looks simpler and faster?

Sorry for slow activity here, got a little bit taken back by all JIT-errors.

The reason for choosing b)

    uniforms = torch.rand_like(logits)
    uniforms = torch.distributions.utils.clamp_probs(uniforms)
    gumbels = -((-(uniforms.log())).log())  # -log(~Exp(1)) = ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)

Instead of c)

    gumbels = -logits.new(logits.shape).exponential_().log() # ~Gumbel(0,1)
    gumbels = (logits + gumbels) / tau  # ~Gumbel(logits,tau)
    y_soft = gumbels.softmax(dim)  # Gumbel softmax noise

Was because alicanb had the good point that implementation be similar to RelaxedOneHotCategorical and b) generates identical noise to RelaxedOneHotCategorical. Now that we know it wont work with JIT maybe b) is not worth the hassle. I just need to verify first that .exponential_() will work with JIT too, in that case that is preferable