nnf_gumbel_softmax: Gumbel_softmax

View source: R/nnf-activation.R

nnf_gumbel_softmaxR Documentation

Gumbel_softmax

Description

Samples from the Gumbel-Softmax distribution and optionally discretizes.

Usage

nnf_gumbel_softmax(logits, tau = 1, hard = FALSE, dim = -1)

Arguments

logits

⁠[..., num_features]⁠ unnormalized log probabilities

tau

non-negative scalar temperature

hard

if True, the returned samples will be discretized as one-hot vectors, but will be differentiated as if it is the soft sample in autograd

dim

(int) A dimension along which softmax will be computed. Default: -1.


torch documentation built on May 29, 2024, 9:54 a.m.