import torch


def init_mask(aux):
    """
    Initializes masks of weight to freeze

    Parameters
    ----------
    aux: dict
        A map from weight name to tuple (W, Z, U, project_fun)

    Returns
    -------
    mask: dict
        A map from weight name to tuple (W, mask)
    """
    with torch.no_grad():
        mask = {}
        for weight_name, (W, _, _, project_fun) in aux.items():
            _, m = project_fun(W)
            W[m] = 0
            mask[weight_name] = (W, m)
    return mask