import torch


def concrete_binarize(x, t=0.1, u=None):
    """
    Binarize function from the paper
    'DropMax: Adaptive Variational Softmax'
    http://papers.nips.cc/paper/7371-dropmax-adaptive-variational-softmax
    Args:
        x (torch.FloatTensor): logits to binarize
        t (float): temperature
        u (torch.FloatTensor): uniformly distributed noise to mix-up with logits
    """
    u = torch.rand(*x.shape).to(x.device) if u is None else u
    eps = 1e-7
    return torch.sigmoid(1/t * (x + torch.log(u/(1-u) + eps)))