2016 paper at: https://arxiv.org/abs/1611.01144
tfp has a ExpRelaxedOneHotCategorical but the temperature must be a float. The Exp part of the name refers to another reparametrization that avoids underflow in the logp.
As the temperature \tau approaches 0, the gumbel softmax version becomes a categorical distribution.
I don’t get why their probability density is well-defined because some y_{i} can be 0. Then I get division by 0 in the \pi_{i} / y_{i}^{\tau+1} terms.
For a categorical variable Y with k = 3 classes, the 3 - 1 dimensional simplex is an ordinary triangle.
Let \vec{y} \in \Delta^{k} be one-hot encoding of the categorical variable with k classes, where \Delta^{k-1} is a (k - 1) dimensional probability simplex.
Let \pi_{i} = P(Y = i) be the probability that the categorical random variable have a value of i.
P_{\pi, \tau}(\vec{y}) = \Gamma(k) \tau^{k-1} \left( \sum_{i=1}^{k} \pi_{i} / y_{i}^{\tau} \right)\prod_{i=1}^{k} (\pi_{i} / y_{i}^{\tau+1})
I guess \Gamma(k) means factorial of k when k is an integer.