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

        # 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

    Compute a binary mask for pruning based on magnitude-based ranking
    If granularity is `per_scalar`, L1 norm is used. L2 is used otherwise
    
coremltools/optimize/torch/pruning/_utils.py