code: def test_magnitude_pruning_granularity_parameter_usage(
      simple_module, granularity
  ):
      """
      Tests MagnitudePruner creates mask of the correct shape
      depending on the granularity parameter.
  
      We set target sparsity to 1.0 so the mask should be all zeros after 4 iterations.
      """
      config = MagnitudePrunerConfig.from_dict(
          {
              "global_config": {
                  "scheduler": {"update_steps": [2, 3]},
                  "initial_sparsity": 0.5,
                  "target_sparsity": 1.0,
                  "granularity": granularity,
              }
          },
      )
      pruner = MagnitudePruner(simple_module, config)
      simple_module = pruner.prepare()
  
      # Perform 4 iterations
      for _ in range(4):
          pruner.step()
  
      mask_data = simple_module.conv2d.weight_mask.numpy()
      # Pruning mask should be all zeros since the pruner should be at 100% sparsity.
      if granularity == "per_scalar":
          expected_mask_shape = (4, 3, 3, 3)
      else:
          assert granularity == "per_channel"
          expected_mask_shape = (4, 1, 1, 1)
      np.testing.assert_equal(mask_data, np.zeros(expected_mask_shape))
  
imports: import copy
  from collections import OrderedDict
  import numpy
  import pytest
  import torch
  from coremltools.optimize.torch.pruning import MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig
  from coremltools.optimize.torch.pruning._utils import n_m_mask
label_desc: Use image node_id 16 for calling a global function with example usage: test_magnitude_pruning_granularity_parameter_usage(simple_module, granularity) without return types