We might consider implementing RelaxedOneHotCategorical via a BoltzmannTransform and Gumbel (with the appropriate location), and overriding the log_prob method as described in #123. This is fewer lines of code, but might impact performance, and it's worth investigating briefly.