tensorplay.nn.functional.gumbel_softmax

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

Sample from the Gumbel-Softmax distribution and optionally discretize.

straight-through when hard=True).