code: def n_m_mask(
      weights: _torch.Tensor,
      nm: _Tuple[int, int],
      dim: _Optional[int] = 1,
  ):
      """
      Create a n:m sparsity mask.
      """
      shape = weights.shape
      permuted_shape = shape
      rank = len(shape)
      num_zeros, block_size = nm
      mask_value = 0.0
  
      assert num_zeros < block_size, (
          f"n (number of zeros) = {num_zeros} must be "
          f"less than m (block size) = {block_size}"
      )
  
      assert dim in [0, 1], (
          f"n:m mask is supported along dimensions (0, 1), "
          f"corresponding to input and output channels. Received "
          f"dim = {dim}"
      )
      # rank 2: torch.Linear, rank 3: torch.Conv1d,
      # rank 4: torch.Conv2d, rank 5: torch.Conv3d
      assert rank in [
          2,
          3,
          4,
          5,
      ], f"weights tensor rank must be in [2, 3, 4, 5], got {rank}"
  
      # num_non_zeros = block_size - num_zeros
  
      # if n:m is required along C_o, flip C_i and C_o
      if dim == 0:
          weights = _torch.permute(
              weights, [1, 0] + list(range(2, rank))
          )
      # transform to A x C_i
      # For Conv1D: C_o x C_i x H         ==>         H x C_o x C_i ==> H*C_o     x C_i
      # For Conv2D: C_o x C_i x H x W     ==>     H x W x C_o x C_i ==> H*W*C_o   x C_i
      # For Conv3D: C_o x C_i x H x W x D ==> H x W x D x C_o x C_i ==> H*W*D*C_o x C_i
      if rank > 2:
          permute_array = list(range(2, rank)) + [0, 1]
          weights = _torch.permute(weights, permute_array)
          permuted_shape = weights.shape
          weights = _torch.reshape(weights, (-1, weights.shape[-1]))
  
      abs_weights = weights.abs()
      padding_size = block_size - abs_weights.shape[-1] % block_size
      abs_weights_pad = _torch.nn.functional.pad(
          abs_weights, (0, padding_size), mode="constant"
      )
  
      num_blocks = abs_weights_pad.numel() // block_size
      weights_blocks = abs_weights_pad.view(num_blocks, block_size)
  
      indices = _torch.argsort(weights_blocks, dim=1)[:, :num_zeros]
      sparsity_mask = _torch.ones(
          [num_blocks, block_size], device=weights.device
      )
      sparsity_mask.scatter_(dim=1, index=indices, value=mask_value)
      sparsity_mask = sparsity_mask.view(abs_weights_pad.shape)
      sparsity_mask = sparsity_mask[:, : abs_weights.shape[-1]]
  
      # revert changes to mask shape to achieve same size as original weight
      if rank > 2:
          sparsity_mask = _torch.reshape(sparsity_mask, permuted_shape)
          permute_array = [rank - 2, rank - 1] + list(
              range(0, rank - 2)
          )
          sparsity_mask = _torch.permute(sparsity_mask, permute_array)
      if dim == 0:
          sparsity_mask = _torch.permute(
              sparsity_mask, [1, 0] + list(range(2, rank))
          )
  
      return sparsity_mask
  
file_path: coremltools/coremltools/optimize/torch/pruning/_utils.py
imports: import logging
  from typing import List
  from typing import Optional
  from typing import Tuple
  from typing import cast
  import numpy
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: n_m_mask(weights, nm, dim) and returns: sparsity_mask