feat: add gumbel_softmax

[OrangeX4]

Code of RelaxedOneHotCategorical came from:

from torch.distributions import RelaxedOneHotCategorical

We inherit OneHotCategorical to use dist.kl_divergence(), and the torch.distributions.RelaxedOneHotCategorical does not implement dist.kl_divergence().