Magnitude-Based Pruning, Single Step
The problem, from first principles
Quantization (02-quantization) shrinks a model by storing each weight with fewer bits. Pruning takes a different approach: instead of representing every weight more cheaply, remove some weights entirely by setting them to exactly zero — a matrix with many exact zeros compresses well and, with the right hardware/software support, skips computation on those zeros altogether.
From theory to code
Implement magnitude_prune(weight, sparsity), which zeros out the smallest-magnitude fraction of a weight array's elements, leaving the rest untouched. The underlying premise: a weight already close to zero contributes almost nothing to the layer's output, so removing the smallest ones should hurt accuracy far less than removing an equally-sized random selection would.
Constraints
weight: any-shape float array.sparsity: a fraction in[0, 1], the proportion of elements to zero out.- Elements are ranked by absolute value; the
sparsityfraction with the smallest magnitude get zeroed. sparsity <= 0prunes nothing (returns an unchanged copy).- Returns a new array — never modifies
weightin place.
Hints
Open one at a time. Each gives away a little more than the last.
Hint 1
np.partition(flat_abs, k-1)[k-1] finds the value that would sit at position k-1 in a full sort, without actually sorting the whole array — exactly the threshold value at "the smallest k elements' boundary."
Hint 2
Once you have that threshold, a boolean mask np.abs(weight) > threshold is True for every element that survives pruning — multiplying weight by that mask zeros out everything else in one step.
Keywords
No results yet
Run your code to check it against the examples, or submit to run every test.