def cond_gumbel_sample(all_joint_log_probs,
perturbed_log_probs) -> torch.Tensor:
# Sample plates x k? x |D_yv| Gumbel variables
gumbel_d = Gumbel(loc=all_joint_log_probs, scale=1.0)
G_yv = gumbel_d.rsample()
# Condition the Gumbel samples on the maximum of previous samples
# plates x k
Z = G_yv.max(dim=-1)[0]
T = perturbed_log_probs
vi = T - G_yv + log1mexp(G_yv - Z.unsqueeze(-1))
# plates (x k) x |D_yv|
return T - vi.relu() - torch.nn.Softplus()(-vi.abs())
def rsample_gumbel(
distr: Distribution,
n: int,
) -> torch.Tensor:
gumbel_distr = Gumbel(distr.logits, 1)
return gumbel_distr.rsample((n, ))