code: def magnitude_ranked_mask(
      weights: _torch.Tensor,
      sparsity_fraction: float,
      block_size: int,
      granularity: str,
  ) -> _torch.Tensor:
      """
      Compute a binary mask for pruning based on magnitude-based ranking
      If granularity is `per_scalar`, L1 norm is used. L2 is used otherwise
      """
      shape = weights.shape
      rank = len(shape)
  
      # rank 1: flattened global unstructured weights, rank 2: torch.Linear, rank 3: torch.Conv1d,
      # rank 4: torch.Conv2d, rank 5: torch.Conv3d
      assert rank in [
          1,
          2,
          3,
          4,
          5,
      ], f"weights tensor rank must be in [1, 2, 3, 4, 5], got {rank}"
  
      if granularity == "per_scalar" or rank == 2:
          magnitude_map = weights.abs()
          nb_weight_components = weights.numel()
  
      elif rank in [3, 4, 5]:
          if granularity == "per_kernel":
              start_dim = 2
              nb_weight_components = shape[0] * shape[1]
          elif granularity == "per_channel":
              start_dim = 1
              nb_weight_components = shape[0]
          else:
              raise ValueError(
                  f"Unsupported granularity for magnitude_ranked_mask: {granularity}"
              )
  
          # Compute L2 norm per weight slice (as defined by the granularity)
          magnitude_map = _torch.norm(
              weights.flatten(start_dim), dim=-1
          )
          for _ in range(rank - start_dim):
              magnitude_map = magnitude_map.unsqueeze(-1)
  
      if block_size > 1:
          ch_shape = shape[0]
          if ch_shape % block_size != 0:
              # Since the number of channels isn't divisible by block size,
              # we shall pad the channels so that it is divisible
              pad_shape = list(magnitude_map.shape)
              pad_shape[0] = block_size - ch_shape % block_size
              magnitude_map = _torch.cat(
                  [
                      magnitude_map,
                      _torch.zeros(
                          pad_shape, device=magnitude_map.device
                      ),
                  ],
                  dim=0,
              )
              ch_shape = magnitude_map.shape[0]
              assert ch_shape % block_size == 0
  
          if block_size > ch_shape / 2:
              raise ValueError(
                  f"Pruning block size ({block_size}) can be at most half the number of output channels ({ch_shape}/2={ch_shape/2})"
              )
  
          # Reshape to expose the "block" sub-axis
          s = list(magnitude_map.shape)  # block exposed shape
          s.insert(1, block_size)
          s[0] = int(s[0] / block_size)
          f = [-1] * len(s)  # expand factors to recover orig shape
          f[1] = block_size
          magnitude_map = (
              magnitude_map.view(s)
              .pow(2)
              .sum(1, keepdim=True)
              .sqrt()
              .expand(f)
              .contiguous()
              .view(magnitude_map.shape)
          )
  
          # Reshape to original shape in case of padding
          magnitude_map = magnitude_map[: shape[0]]
  
      nb_nonzero = _torch.ceil(
          _torch.as_tensor(nb_weight_components, dtype=_torch.float32)
          * (1 - sparsity_fraction)
      ).int()
  
      # handle special case when sparsity_fraction = 1.0
      if nb_nonzero == 0:
          thr = 1.0 + magnitude_map.flatten().max()
      else:
          thr = (
              magnitude_map.flatten().sort()[0].flip(0)[nb_nonzero - 1]
          )  # produces same mask for 1.0 and 0.0 sparsity
  
      mask = _torch.greater_equal(magnitude_map, thr)
  
      return 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 3 for calling a global function with example usage: magnitude_ranked_mask(weights, sparsity_fraction, block_size, granularity) and returns: mask