Intro I recently learned about the so-called Gumbel-Max and Gumbel-Softmax tricks. Essentially, the Gumbel-Max trick says that if we have a categorical distribution $\vec{\pi} = {\pi_1, \ldots \pi_K}$ and i.i.d. $\mathrm{Gumbel}(0, 1)$-distributed random variables $G_i, \; 1\le i\le K$, then $$ \forall k \quad \mathbb{P}(G_k + \log(\pi_k) = \max\{G_i + \log(\pi_i) \colon 1 \le i \le K\}) = \pi_k.