Yes I did read that paper and the Gumbel Softmax one as well. I was just confused about the way log_prob is used. I see that the KL can be defined in two ways:

Using the parameterized distribution where the samples are sampled from the q distribution.

kl = log_q_y - log_p_y


Using the logits directly

postlogprob = log(softmax(logits))
priorlogprob = log(softmax(prior))
kl= (poslogprob * samples) - (priorlogprob * samples)

Am I correct here?