code: def test_magnitude_pruning_correctness(simple_module):
      """
      Test correctness of magnitude pruning.
  
      Initialize convolution weight with 4 output channels,
      with weights associated with channel `k` having integer value k+1 (k=0,...,3).
      We test that pruning twice indeed zeros out 3 output channels.
      """
      config = MagnitudePrunerConfig.from_dict(
          {
              "global_config": {
                  "scheduler": {"update_steps": [2, 3]},
                  "initial_sparsity": 0.0,
                  "target_sparsity": 0.75,
                  "granularity": "per_channel",
              }
          },
      )
      pruner = MagnitudePruner(simple_module, config)
      simple_module = pruner.prepare()
  
      # Perform 4 iterations: pruning should happen on steps 2 and 3
      # step 1: No pruning
      pruner.step()
      np.testing.assert_equal(
          simple_module.conv2d.weight_mask.numpy(),
          np.array([1, 1, 1, 1], dtype=np.int32).reshape((4, 1, 1, 1)),
      )
      # step 2: prune once, polynomial schedule will give new sparsity as 0.0, still no pruning
      pruner.step()
      np.testing.assert_equal(
          simple_module.conv2d.weight_mask.numpy(),
          np.array([1, 1, 1, 1], dtype=np.int32).reshape((4, 1, 1, 1)),
      )
      # step 3: prune once again, polynomial schedule will give new sparsity as 1.0, 75% = 3 out of 4
      # channels with least magnitude (first three channels) will be pruned out
      pruner.step()
      np.testing.assert_equal(
          simple_module.conv2d.weight_mask.numpy(),
          np.array([0, 0, 0, 1], dtype=np.int32).reshape((4, 1, 1, 1)),
      )
  
      # step 4: prune once again, polynomial schedule sparsity stays at 0.75, no further pruning
      pruner.step()
      np.testing.assert_equal(
          simple_module.conv2d.weight_mask.numpy(),
          np.array([0, 0, 0, 1], dtype=np.int32).reshape((4, 1, 1, 1)),
      )
  
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 14 for calling a global function with example usage: test_magnitude_pruning_correctness(simple_module) without return types