Add Axon.Pruning with magnitude-based unstructured pruning

[seanmor5]

Closes #120.

Pruning is a standard compression step alongside quantization, but Axon had nothing for it: there was no way to zero out low-magnitude weights, no way to keep them at zero while fine-tuning, and no way to report how sparse a model state is. This PR adds Axon.Pruning, a sibling of Axon.Quantization that works purely on an Axon.ModelState and leaves the model graph untouched.

Design

This is a first, deliberately narrow increment: unstructured magnitude pruning. The parameters with the smallest absolute values are set to zero, either against one global threshold (scope: :global, the default, so some tensors end up sparser than others) or per tensor (scope: :tensor, every selected tensor is pruned to exactly the requested sparsity). Pruning always produces a mask alongside the pruned state: a nested map mirroring the model state data with a {:u, 8} tensor (1 = keep, 0 = pruned) for every pruned parameter. Parameters that were not selected are simply absent from the mask.

The public API is:

• magnitude_mask/3 computes the mask without touching the state. Candidates come from the parameters tree only, so batch-norm running statistics (which live in state) are never pruned. Tied parameters (Axon.ModelState.SharedParameter), quantized parameters (Axon.Quantization.QTensor) and non-float tensors are skipped. The :filter option receives the same access paths as Axon.ModelState.freeze/2 (e.g. ["dense_0", "kernel"] or ["lstm_0", "input_kernel", "wii"]); the default keeps any path with a name ending in "kernel", which covers dense/conv/embedding kernels and the composite RNN kernels while leaving biases and normalization parameters alone.
• apply_mask/2 zeroes the masked entries of a model state or a plain parameter map using Nx.select/3, so dtypes are preserved and NaN/Inf don't leak through a multiply.
• prune/3 is magnitude_mask/3 + apply_mask/2 and returns {pruned_state, mask}.
• masked_optimizer/2 wraps anything Axon.Loop.trainer/4 accepts (a Polaris.Optimizers atom or an {init_fn, update_fn} tuple) via Polaris.Updates.stateful/3. The mask is stored in the optimizer state, and the final update is zeroed at pruned positions, so a weight that starts at zero stays exactly zero regardless of momentum or weight decay. This composes with the existing loop and needs no changes to Axon.Loop. Mask entries for parameters that are not trained (e.g. frozen ones) are ignored.
• sparsity/1 and global_sparsity/1 report the fraction of exactly-zero entries per parameter (nested maps for composite parameters) and size-weighted across the whole state.

The mask computation ranks magnitudes with a double Nx.argsort (stable), so exactly round(sparsity * n) entries are pruned even when magnitudes tie, and it works inside defn on any backend. magnitude_mask/3, apply_mask/2 and prune/3 are deftransforms so they can be called from defn; the sparsity reporting functions call Nx.to_number/1 and are eager only.

Usage

{pruned_state, mask} = Axon.Pruning.prune(model_state, 0.9)

Axon.Pruning.global_sparsity(pruned_state)
#=> 0.87...

trained_state =
  model
  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Pruning.masked_optimizer(:adam, mask))
  |> Axon.Loop.run(data, pruned_state, epochs: 2)

Limitations

• Zeroed weights are still stored densely: this reduces the number of non-zero parameters, not memory or latency. Getting a speedup needs sparse tensor support or structured pruning, neither of which exists yet.
• Structured/channel pruning, gradual pruning schedules and rewriting the model graph are out of scope for this increment.
• The default filter is a naming heuristic on "kernel"; anything else needs an explicit :filter.
• masked_optimizer/2 copies the mask to the binary backend so it can be embedded in the traced optimizer init; from then on it travels as ordinary optimizer state (and is therefore checkpointed with it).
• Ties between equal magnitudes are broken by position, deterministically but arbitrarily.

Tests

test/axon/pruning_test.exs covers per-tensor and global masks against hand-computed expectations, exact counts under ties, the sparsity 0/1 short-circuits, the default filter on dense + batch norm (state untouched) and LSTM (composite kernels masked, biases not), custom filters, skipping of tied, quantized and integer parameters, argument validation, apply_mask/2 on model states and plain parameter maps (ignoring mask keys that are absent), all three transforms inside defn, sparsity/1/global_sparsity/1 including composite parameters and empty states, and masked_optimizer/2 with both :adam and an {init_fn, update_fn} tuple through a real Axon.Loop run (pruned entries stay exactly zero, unmasked entries train, and a control run without the wrapper does not keep them at zero). The suite passes on the default backend and with USE_EXLA=1.

🤖 Generated with Claude Code

[polvalente]

Suggested change

	{optimizer, 0} in Polaris.Optimizers.__info__(:functions) do
	function_exported?(Polaris.Optimizers, optimizer, 0) do

[polvalente]

Suggested change

	tensor |> Nx.equal(0) |> Nx.sum() |> Nx.to_number()
	tensor |> Nx.equal(0) |> Nx.as_type(:u64) |> Nx.sum() |> Nx.to_number()

Otherwise there will be overflows happening