API Reference

This is the exhaustive, auto-generated reference for every public module of ShrinkAI, browsable via the navigation tree on the left, drill into a package to see its modules, and into a module to see its classes and functions, complete with signatures and source.

There is no flat import shrinkai API: everything is imported from its submodule, e.g. from shrinkai.distillation import Distiller. If you're looking for a guided introduction instead of a reference, start with the Quickstart then the Tutorials.

Where to start, by task

I want to...	Start here
Distill a teacher into a smaller student	Distiller
Pick or write a distillation loss	shrinkai.distillation.losses
Stop training early / save the best checkpoint	EarlyStopping, ModelCheckpoint
Compare intermediate representations (teacher vs. student)	FeatureExtractor, FeatureAnalyzer
Reconcile mismatched feature dimensions	FeatureProjector, AttentionHeadSelector
Prune a model	Pruner (masking), ChannelPruner (physical shrink)
Quantize a model (PTQ / QAT)	Quantizer
Measure latency, size, params, FLOPs	Profiler, count_flops
Export a trained model for deployment	export_onnx, export_torchscript

Package map

• shrinkai.distillation

  Train a student to mimic a teacher: Distiller (high-level facade), DistillationEngine (training loop), ready-made callbacks, and a library of losses under shrinkai.distillation.losses.

• shrinkai.adapters

  Bridge teacher/student architectures: hook-based feature extraction and dimension-matching projectors.

• shrinkai.analysis

  Score how well a student's internal representations align with its teacher's (CKA, RSA, spatial attention).

• shrinkai.compression

  Shrink a model after (or during) training: pruning (shrinkai.compression.pruning) and quantization (shrinkai.compression.quantization).

• shrinkai.profiler

  Measure and compare parameter count, disk size, latency, FLOPs, and memory footprint.

• shrinkai.export

  Get a trained/compressed model out of Python: ONNX and TorchScript export.