Bryor Snefjella

> I believe https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/RelaxedOneHotCategorical is what you are looking for (the distribution goes under the name Relaxed One Hot Categorical, Gumbel Softmax and Concrete in the literature). Yeah, the distribution...

From the looks of https://github.com/tensorflow/probability/blob/v0.17.0/tensorflow_probability/python/distributions/mixture_same_family.py#L266-L270 the changes needed to allow this don't look too complicated. In _sample_n, the mask can be replaced with samples from the relaxed_one_hot_categorical mixture selecting distribution...

My network gives nans whenever I've tried setting that to true (IWAE with a mixture proposal distribution). I was able to set up a gumbel mixture, no nans. Didn't look...