def test_constant_sparsity_correctness(simple_module, step_and_target):
    """
    Tests correctness of spline schedule.

    Note: Schedule can be stepped beyond the maximum step specified in update_steps.
    Beyond the max update step, the sparsity stays at the target sparsity.
    For example, in the first test case, we step 10 times, whereas max step is 9. At the 10th call
    to schedule.step, sparsity remains at 1.0.
    """
    begin_step, target_sparsity = step_and_target
    initial_sparsity = target_sparsity if begin_step == 0 else 0.0
    config = MagnitudePrunerConfig().set_global(
        ModuleMagnitudePrunerConfig(
            scheduler=ConstantSparsityScheduler(begin_step=begin_step),
            initial_sparsity=initial_sparsity, target_sparsity=target_sparsity,
        )
    )
    pruner = MagnitudePruner(simple_module, config)
    pruner.prepare(inplace=True)
    for _ in range(begin_step):
        assert pruner._pruner_info[''].sparsity_level == initial_sparsity
        pruner.step()
    assert pruner._pruner_info[''].sparsity_level == target_sparsity

    Tests correctness of spline schedule.

    Note: Schedule can be stepped beyond the maximum step specified in update_steps.
    Beyond the max update step, the sparsity stays at the target sparsity.
    For example, in the first test case, we step 10 times, whereas max step is 9. At the 10th call
    to schedule.step, sparsity remains at 1.0.
    
coremltools/test/optimize/torch/pruning/test_pruning_scheduler.py