Function `gumbel_softmax`'s documentation is misleading

Boltzmachine: <p>In the documentation of <code>gumbel_softmax</code>, the first parameter <code>logits</code></p>
<pre><code class="lang-auto">logits: `[..., num_features]` unnormalized log probabilities
</code></pre>
<p>It confused me a lot that why the logtis could be unnormalized. In the equation the torch function uses</p>
<pre><code class="lang-auto">softmax((log p_i - log (-log e_i)) / t)
</code></pre>
<p>where <code>log(-log e_i))</code> is the gumbel noise, <code>t</code> is the temperature, <code>p_i</code> is the probability. if you want to use the unnormalized logits <code>o_i</code> to replace <code>log p_i</code> in the equation may not hold unless you have</p>
<pre><code class="lang-auto">\sum_i exp(o_i) = 1
</code></pre>
<p>which is not guaranteed by an arbitrary network.</p>
<p>This <a href="https://discuss.pytorch.org/t/gumbel-softmax-function-logits/22797">thread</a> also mentioned this issue</p>

KFrank: <p>Hi Weikang!</p>
<p>I do not have an opinion about whether or not pytorch’s <code>gumbel_softmax()</code><br>
is correct.</p>
<aside class="quote no-group" data-username="Boltzmachine" data-post="1" data-topic="186157" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/boltzmachine/48/32323_2.png" class="avatar"> Boltzmachine:</div>
<blockquote>
<p>the equation may not hold unless you have</p>
<pre><code class="lang-auto">\sum_i exp(o_i) = 1
</code></pre>
<p>which is not guaranteed by an arbitrary network.</p>
</blockquote>
</aside>
<p>However, as a practical matter, if you are concerned about this issue, you<br>
may pass the output of your “arbitrary network” through <code>log_softmax()</code><br>
before passing it to <code>gumbel_softmax()</code>.</p>
<p>One way of looking at this is that the only thing that <code>log_softmax()</code> does<br>
is convert unnormalized log-probabilities to normalized log-probabilities.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.1
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; unnormalized_log_probabilities = torch.randn (5)
&gt;&gt;&gt; unnormalized_log_probabilities
tensor([-1.2075,  0.5493, -0.3856,  0.6910, -0.7424])
&gt;&gt;&gt;
&gt;&gt;&gt; probabilities = unnormalized_log_probabilities.softmax (0)
&gt;&gt;&gt; probabilities
tensor([0.0577, 0.3342, 0.1312, 0.3851, 0.0918])
&gt;&gt;&gt; probabilities.sum()
tensor(1.)
&gt;&gt;&gt;
&gt;&gt;&gt; normalized_log_probabilities = unnormalized_log_probabilities.log_softmax (0)
&gt;&gt;&gt; normalized_log_probabilities
tensor([-2.8528, -1.0960, -2.0309, -0.9544, -2.3877])
&gt;&gt;&gt; normalized_log_probabilities.exp()
tensor([0.0577, 0.3342, 0.1312, 0.3851, 0.0918])
&gt;&gt;&gt; normalized_log_probabilities.exp().sum()
tensor(1.)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (normalized_log_probabilities.log_softmax (0), normalized_log_probabilities)
True
</code></pre>
<p>(Closely related is that if you start with <em>normalized</em> log-probabilities, then<br>
the inverse of <code>softmax()</code> is <code>log()</code>.)</p>
<p>As an aside, log-probabilities are not <em>logits</em> (and pytorch should not use<br>
<code>logits</code> as the name of the log-probabilities argument to <code>gumbel_softmax()</code>).<br>
I’ve mixed up the terminology in some of my posts, but it’s worth trying to get<br>
it right.  <em>Logits</em> are log-odds-ratios – <code>log (p / (1 - p))</code> (and are converted<br>
to probabilities with the logistic <code>sigmoid()</code> function).  Logits play the same<br>
role with <code>BCEWtthLogitsLoss</code> that (unnormalized) log- probabilities play<br>
with <code>CrossEntropyLoss</code>.</p>
<p>Best.</p>
<p>K. Frank</p>