Stretch: Iterative Pruning Schedule

The problem, from first principles

01-magnitude-pruning zeros out a fixed fraction of weights in a single shot. Pruning a large fraction all at once tends to genuinely hurt accuracy — the network hasn't had any chance to adapt its remaining weights to compensate for what was just removed. Spreading the same total amount of pruning across many smaller steps, interleaved with retraining between them in a real pipeline, tends to reach the same final sparsity with far less accuracy loss.

From theory to code

Implement cubic_sparsity_schedule(step, total_steps, target_sparsity), which computes how much sparsity to reach by a given step, and iterative_prune(weight, target_sparsity, num_steps), which applies 01-magnitude-pruning's magnitude_prune repeatedly, once per step, following that schedule.

Constraints

• cubic_sparsity_schedule(0, ...) returns 0.0; cubic_sparsity_schedule(total_steps, ...) returns target_sparsity exactly.
• The schedule is monotonically increasing and front-loaded (covers most of the sparsity increase in the earlier steps, tapering off as it approaches the target).
• A step beyond total_steps clamps to target_sparsity.
• iterative_prune returns a list of num_steps arrays, each pruned to that step's scheduled sparsity, with sparsity never decreasing from one step to the next.

Hints

Open one at a time. Each gives away a little more than the last.