def test_magnitude_pruner_block_sparsity(out_channels, block_size):
    """
    Test block sparsity structure is obtained by MagnitudePruner when block_size > 1
    """
    conv2d = torch.nn.Conv2d(in_channels=3,
                             out_channels=out_channels,
                             kernel_size=(3, 3),
                             bias=False,
                             groups=1)
    weight_shape = tuple(conv2d.weight.shape)
    weight_tensor = torch.abs(torch.randn(*weight_shape))
    weight_tensor[weight_tensor == 0] = 1.0
    conv2d.weight = torch.nn.Parameter(weight_tensor)
    config = MagnitudePrunerConfig.from_dict(
            {"global_config":
                {
                    "scheduler": {"update_steps": [1, 2]},
                    "initial_sparsity": 0.0,
                    "target_sparsity": 0.5,
                    "block_size": block_size,
                }},
        )
    pruner = MagnitudePruner(conv2d, config)
    conv2d = pruner.prepare()

    for _ in range(4):
        pruner.step()

    if block_size > 1:
        for idx in range(0, out_channels, block_size):
            for jdx in range(1, block_size):
                assert torch.all(conv2d.weight_mask[idx] == conv2d.weight_mask[idx + jdx])

    assert torch.sum(conv2d.weight_mask == 0).item() == int(0.5 * torch.numel(conv2d.weight))

    Test block sparsity structure is obtained by MagnitudePruner when block_size > 1
    
coremltools/test/optimize/torch/pruning/test_magnitude_pruner.py