Module pruning 

Neural network pruning: importance scoring, sparsity masks, and compression. Neural network pruning module.

This module provides importance scoring and sparsity mask generation for neural network pruning, following the specification in docs/specifications/advanced-pruning.md.

Toyota Way Principles

• Jidoka: All numerical operations validate for NaN/Inf
• Poka-Yoke: Type-safe patterns prevent invalid configurations
• Genchi Genbutsu: Calibration uses real activation data

Example

use aprender::pruning::{MagnitudeImportance, Importance};
use aprender::nn::Linear;

let layer = Linear::new(512, 256);
let importance = MagnitudeImportance::l2();
let scores = importance.compute(&layer, None).expect("valid layer weights");

println!("Importance stats: min={}, max={}", scores.stats.min, scores.stats.max);

References

• Han, S., et al. (2015). Learning both weights and connections. NeurIPS.
• Sun, M., et al. (2023). A simple and effective pruning approach. arXiv:2306.11695.
• Frantar, E., & Alistarh, D. (2023). SparseGPT. ICML.
• Frankle, J., & Carbin, M. (2018). The Lottery Ticket Hypothesis. arXiv:1803.03635.

Structs

ActivationStats

Per-layer activation statistics.

BlockImportanceScores

Block Importance (BI) scores for all layers.

BlockSparseTensor

Block sparse tensor representation.

COOTensor

Coordinate (COO) sparse tensor representation.

CSRTensor

Compressed Sparse Row (CSR) representation.

CalibrationContext

Calibration context holding activation statistics for all layers.

ChannelImportance

Channel importance scores.

DependencyGraph

Dependency graph for a neural network.

DepthPruner

Minitron Depth Pruner: Layer removal based on Block Importance.

DepthPruningResult

Result of depth (layer) pruning operation.

GraphEdge

An edge in the dependency graph.

GraphNode

A node in the dependency graph representing a layer or operation.

ImportanceScores

Importance scores with metadata.

ImportanceStats

Statistical summary of importance scores.

LotteryTicketConfig

Configuration for Lottery Ticket pruning.

LotteryTicketPruner

Lottery Ticket Hypothesis pruner.

LotteryTicketPrunerBuilder

Builder for LotteryTicketPruner.

MagnitudeImportance

Magnitude-based importance estimator.

MagnitudePruner

Simple magnitude-based pruner.

PruningPlan

Pruning plan that ensures consistency across dependent layers.

PruningResult

Result of a pruning operation with diagnostics.

SparseGPTImportance

SparseGPT importance estimator using Hessian-based saliency.

SparsityMask

Sparsity mask with validation.

WandaImportance

Wanda (Weights and Activations) importance estimator.

WandaPruner

Wanda-based pruner.

WidthPruner

Minitron Width Pruner: Channel removal based on activation importance.

WidthPruningResult

Result of width (channel) pruning operation.

WinningTicket

A “winning ticket” - the sparse subnetwork found by LTH.

Enums

DependencyType

Type of dependency between layers.

NodeType

Type of layer/operation node.

NormType

Norm type for magnitude computation.

PruningError

Pruning operation errors with detailed context.

RewindStrategy

Strategy for rewinding weights after pruning.

SparseFormat

Sparse tensor format enumeration.

SparseTensor

Unified sparse tensor type.

SparsityPattern

Sparsity pattern constraints.

Traits

Importance

Core trait for importance estimation algorithms.

Pruner

High-level pruning interface.

Functions

generate_block_mask

Generate a block sparsity mask.

generate_column_mask

Generate a column sparsity mask.

generate_nm_mask

Generate an N:M structured sparsity mask.

generate_row_mask

Generate a row sparsity mask.

generate_unstructured_mask

Generate an unstructured sparsity mask based on importance scores.

propagate_channel_pruning

Propagate channel pruning through a dependency graph.

prune_module

Convenience function to prune a module with a single call.

sparsify

Apply a sparsity mask to a tensor and return sparse representation.