(tensor, threshold, align=None)
Computes a simple binary mask for given magnitude threshold.

    :param tensor: PyTorch tensor
    :type tensor: `torch.Tensor`
    :param threshold: magnitude threshold for pruning
    :type threshold: `float`
    :return: Mask
    :rtype: `torch.Tensor`
Computes a simple binary mask for given magnitude threshold.
Computes a simple binary mask for given magnitude threshold .
def simple_mask(tensor, threshold, align=None):
    """
    Computes a simple binary mask for given magnitude threshold.

    :param tensor: PyTorch tensor
    :type tensor: `torch.Tensor`
    :param threshold: magnitude threshold for pruning
    :type threshold: `float`
    :return: Mask
    :rtype: `torch.Tensor`
    """
    assert tensor.dim() == 1
    if align is None:
        return torch.ge(tensor.abs(), threshold)
    else:
        size = tensor.size(0)
        if size < align:
            raise RuntimeError('Tensor too small for given alignment')
        t = tensor.abs()
        nnz = torch.ge(t, threshold).nonzero().size(0)
        nnz = int(nnz / align) * align
        _, indices = torch.topk(t, nnz)
        ones = torch.ones(nnz,
                          dtype=tensor.dtype,
                          layout=tensor.layout,
                          device=tensor.device)
        mask = torch.zeros_like(tensor).scatter_(0, indices, ones)
        return mask
def simple_mask ( tensor , threshold , align = None ) : assert tensor . dim ( ) == 1 if align is None : return torch . ge ( tensor . abs ( ) , threshold ) else : size = tensor . size ( 0 ) if size < align : raise RuntimeError ( 'Tensor too small for given alignment' ) t = tensor . abs ( ) nnz = torch . ge ( t , threshold ) . nonzero ( ) . size ( 0 ) nnz = int ( nnz / align ) * align _ , indices = torch . topk ( t , nnz ) ones = torch . ones ( nnz , dtype = tensor . dtype , layout = tensor . layout , device = tensor . device ) mask = torch . zeros_like ( tensor ) . scatter_ ( 0 , indices , ones ) return mask