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

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