import torch


def emb_2d_dropout(training, mask_perc, tensor):
    """2D dropout of tensor where entire row gets zero-ed out with prob
    mask_perc.

    Args:
        training: if the model is in train mode or not
        mask_perc: percent to dropout
        tensor: tensor to dropout

    Returns: tensor after 2D dropout
    """
    batch, M, K, dim = tensor.shape
    if training and mask_perc > 0:
        # reshape for masking
        tensor = tensor.contiguous().reshape(batch * M * K, dim)
        # randomly mask each entity embedding
        bern_prob = (torch.ones(batch * M * K, 1) * mask_perc).to(tensor.device)
        zero_mask = torch.bernoulli(bern_prob) > 0
        tensor = tensor.masked_fill(zero_mask, 0)
        tensor = tensor.contiguous().reshape(batch, M, K, dim)
    return tensor