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)))

    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.
    
coremltools/test/optimize/torch/pruning/test_magnitude_pruner.py