(tensor, threshold, blocksize, criteria, align=None)
return mask.reshape(tuple(tmpshape)).permute(tuple(order)).reshape(
        tuple(shape))
Computes an n-D binary mask for given magnitude threshold.

    :param tensor: PyTorch tensor
    :type tensor: `torch.Tensor`
    :param threshold: magnitude threshold for pruning
    :type threshold: `float`
    :param blocksize: desired block size (Tuple)
    :type blocksize: `Tuple`
    :param criteria: aggregation function for thresholding (default: max)
    :type criteria: `condensa.functional`
    :return: Mask
    :rtype: `torch.Tensor`
Computes an n-D binary mask for given magnitude threshold.
Computes an n - D binary mask for given magnitude threshold .
def block_mask(tensor, threshold, blocksize, criteria, align=None):
    """
    Computes an n-D binary mask for given magnitude threshold.

    :param tensor: PyTorch tensor
    :type tensor: `torch.Tensor`
    :param threshold: magnitude threshold for pruning
    :type threshold: `float`
    :param blocksize: desired block size (Tuple)
    :type blocksize: `Tuple`
    :param criteria: aggregation function for thresholding (default: max)
    :type criteria: `condensa.functional`
    :return: Mask
    :rtype: `torch.Tensor`
    """
    # Original implementation at: https://stackoverflow.com/questions/42297115
    # /numpy-split-cube-into-cubes/42298440#42298440
    if tensor.dim() != len(blocksize):
        raise RuntimeError('Tensor and block dimensions do not match')
    ndim = tensor.dim()

    blocksize_flat = np.prod(np.array(blocksize))
    shape = np.array(tensor.shape)
    repeats = (shape / blocksize).astype(int)
    divcheck = (shape % blocksize).astype(int)

    if not np.all(divcheck == 0):
        raise TypeError('Block size must be divisible by tensor size')

    tmpshape = np.column_stack([repeats, blocksize]).ravel()
    order = np.arange(len(tmpshape))
    order = np.concatenate([order[::2], order[1::2]])
    blocks = tensor.abs().reshape(tuple(tmpshape))
    blocks = blocks.permute(tuple(order)).reshape(-1, *blocksize)
    agg = criteria(blocks.reshape(-1, blocksize_flat), dim=1, keepdim=True)

    mask = simple_mask(agg.view(-1), threshold, align)
    mask = mask.view(agg.shape).expand(-1,
                                       blocksize_flat).reshape(blocks.shape)

    N, newshape = mask.shape[0], mask.shape[1:]
    repeats = (shape / newshape).astype(int)
    tmpshape = np.concatenate([repeats, newshape])
    order = np.arange(len(tmpshape)).reshape(2, -1).ravel(order='F')
    return mask.reshape(tuple(tmpshape)).permute(tuple(order)).reshape(
        tuple(shape))
def block_mask ( tensor , threshold , blocksize , criteria , align = None ) : # Original implementation at: https://stackoverflow.com/questions/42297115 # /numpy-split-cube-into-cubes/42298440#42298440 if tensor . dim ( ) != len ( blocksize ) : raise RuntimeError ( 'Tensor and block dimensions do not match' ) ndim = tensor . dim ( ) blocksize_flat = np . prod ( np . array ( blocksize ) ) shape = np . array ( tensor . shape ) repeats = ( shape / blocksize ) . astype ( int ) divcheck = ( shape % blocksize ) . astype ( int ) if not np . all ( divcheck == 0 ) : raise TypeError ( 'Block size must be divisible by tensor size' ) tmpshape = np . column_stack ( [ repeats , blocksize ] ) . ravel ( ) order = np . arange ( len ( tmpshape ) ) order = np . concatenate ( [ order [ : : 2 ] , order [ 1 : : 2 ] ] ) blocks = tensor . abs ( ) . reshape ( tuple ( tmpshape ) ) blocks = blocks . permute ( tuple ( order ) ) . reshape ( - 1 , * blocksize ) agg = criteria ( blocks . reshape ( - 1 , blocksize_flat ) , dim = 1 , keepdim = True ) mask = simple_mask ( agg . view ( - 1 ) , threshold , align ) mask = mask . view ( agg . shape ) . expand ( - 1 , blocksize_flat ) . reshape ( blocks . shape ) N , newshape = mask . shape [ 0 ] , mask . shape [ 1 : ] repeats = ( shape / newshape ) . astype ( int ) tmpshape = np . concatenate ( [ repeats , newshape ] ) order = np . arange ( len ( tmpshape ) ) . reshape ( 2 , - 1 ) . ravel ( order = 'F' ) return mask . reshape ( tuple ( tmpshape ) ) . permute ( tuple ( order ) ) . reshape ( tuple ( shape ) )