def test_polynomial_decay_correctness(simple_module, steps_and_expected):
    """
    Tests correctness of polynomial decay 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.
    """

    update_steps, expected_sparsitys = steps_and_expected
    config = MagnitudePrunerConfig().set_global(
        ModuleMagnitudePrunerConfig(
            scheduler=PolynomialDecayScheduler(update_steps=update_steps),
            initial_sparsity=0.0, target_sparsity=1.0,
        )
    )
    pruner = MagnitudePruner(simple_module, config)
    pruner.prepare(inplace=True)

    for expected in expected_sparsitys:
        pruner.step()
        assert pruner._pruner_info[''].sparsity_level == expected

    Tests correctness of polynomial decay 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