import torch


def get_model_mask_parameters(model, layers):
    """
    Defines a dictionary of type {layer: tensor} containing for each layer of a model, the binary mask representing
    which parameters have a value of zero.
    :param model: PyTorch model.
    :param layers: Tuple of layers on which apply the threshold procedure. e.g. (nn.modules.Conv2d, nn.modules.Linear)
    :return: Mask dictionary.
    """
    mask = {}
    for n_m, mo in model.named_modules():
        if isinstance(mo, layers):
            for n_p, p in mo.named_parameters():
                name = "{}.{}".format(n_m, n_p)
                mask[name] = torch.where(p == 0, torch.zeros_like(p), torch.ones_like(p))

    return mask