tensorplay.nn.functional.gumbel_softmax

tensorplay.nn.functional.gumbel_softmax(logits: TensorBase, tau: float = 1, hard: bool = False, eps: float = 1e-10, dim: int = -1) → TensorBase[source]

Sample from the Gumbel-Softmax distribution and optionally discretize.

Port of torch.nn.functional.gumbel_softmax (reparametrization trick; straight-through when hard=True).