class_docstr: Weight-Norm Sparsifier
  
  This sparsifier computes the norm of every sparse block and "zeroes-out" the
  ones with the lowest norm. The level of sparsity defines how many of the
  blocks is removed.
  
  This sparsifier is controlled by three variables:
  1. `sparsity_level` defines the number of *sparse blocks* that are zeroed-out
  2. `sparse_block_shape` defines the shape of the sparse blocks. Note that
      the sparse blocks originate at the zero-index of the tensor.
  3. `zeros_per_block` is the number of zeros that we are expecting in each
      sparse block. By default we assume that all elements within a block are
      zeroed-out. However, setting this variable sets the target number of
      zeros per block. The zeros within each block are chosen as the *smallest
      absolute values*.
  
  Args:
  
      sparsity_level: The target level of sparsity
      sparse_block_shape: The shape of a sparse block (see note below)
      zeros_per_block: Number of zeros in a sparse block
      norm: Norm to use. Could be either `int` or a callable.
          If `int`, only L1 and L2 are implemented.
  
  Note::
      The `sparse_block_shape` is tuple representing (block_ROWS, block_COLS),
      irrespective of what the rows / cols mean in the data tensor. That means,
      if you were to sparsify a weight tensor in the nn.Linear, which has a
      weight shape `(Cout, Cin)`, the `block_ROWS` would refer to the output
      channels, while the `block_COLS` would refer to the input channels.
  
  Note::
      All arguments to the WeightNormSparsifier constructor are "default"
      arguments and could be overriden by the configuration provided in the
      `prepare` step.
docstr: Creates a block-level mask.
  
  Block-level mask is described as a mask, where the granularity of sparsification of the
  largest patch is the sparse_block_shape. That means that for a given mask and a
  sparse_block_shape, the sparsity is computed only within a patch of a size sparse_block_shape.
  
  In this context the `zeros_per_block` describes the number of zeroed-out elements within a patch.
docstr_tok: ["Creates","a","block-level","mask",".","Block-level","mask","is","described","as","a","mask",",","where","the","granularity","of","sparsification","of","the","largest","patch","is","the","sparse_block_shape",".","That","means","that","for","a","given","mask","and","a","sparse_block_shape",",","the","sparsity","is","computed","only","within","a","patch","of","a","size","sparse_block_shape",".","In","this","context","the","`","zeros_per_block","`","describes","the","number","of","zeroed-out","elements","within","a","patch","."]
code: def _make_block_mask(
      self, data, sparse_block_shape, zeros_per_block, mask=None
  ):
      r"""Creates a block-level mask.
  
      Block-level mask is described as a mask, where the granularity of sparsification of the
      largest patch is the sparse_block_shape. That means that for a given mask and a
      sparse_block_shape, the sparsity is computed only within a patch of a size sparse_block_shape.
  
      In this context the `zeros_per_block` describes the number of zeroed-out elements within a patch.
      """
      h, w = data.shape[-2:]
      block_h, block_w = sparse_block_shape
      dh = (block_h - h % block_h) % block_h
      dw = (block_w - w % block_w) % block_w
      values_per_block = reduce(operator.mul, sparse_block_shape)
  
      if mask is None:
          mask = torch.ones((h + dh, w + dw), device=data.device)
  
      if values_per_block == zeros_per_block:
          # Everything should be sparsified
          mask.data = torch.zeros_like(mask)
          return mask
  
      # create a new padded tensor like data (to match the block_shape)
      padded_data = torch.ones(
          h + dh, w + dw, dtype=data.dtype, device=data.device
      )
      padded_data.fill_(torch.nan)
      padded_data[:h, :w] = data
      unfolded_data = F.unfold(
          padded_data[None, None, :],
          kernel_size=sparse_block_shape,
          stride=sparse_block_shape,
      )
  
      # Temp reshape for mask
      mask_reshape = mask.reshape(unfolded_data.shape)
      _, sorted_idx = torch.topk(
          unfolded_data, k=zeros_per_block, dim=1, largest=False
      )
  
      self._scatter_fold_block_mask(
          dim=1,
          indices=sorted_idx,
          output_shape=padded_data.shape,
          block_shape=sparse_block_shape,
          mask=mask_reshape,
      )
  
      mask.data = (
          mask_reshape.squeeze().reshape(mask.shape).contiguous()
      )
      return mask
  
file_path: pytorch/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py
imports: from functools import reduce
  from typing import Callable, Optional, Tuple, Union
  import torch
  import torch.nn.functional
  from .base_sparsifier import BaseSparsifier
  import operator
label_desc: Use image node_id 4 for calling the WeightNormSparsifier obj's underlying member method code with example usage: obj._make_block_mask(data, sparse_block_shape, zeros_per_block, mask) and returns: mask, mask