Qwix
A Jax library for quantization-aware training and post-training quantization of Flax models, targeting CPUs, GPUs and TPUs as well as on-device LiteRT deployment.
Qwix supports weight-only, dynamic-range and static-range quantization schemes, and three modes: quantization-aware training that emulates quantized behavior while training, post-training quantization tuned for XLA devices such as TPU and GPU, and an ODML mode that annotates a model so LiteRT's converter can produce fully integer models for on-device use. It also supports LoRA and QLoRA.
Numeric formats include native int4, int8 and fp8, and emulated formats from int1 to int7 plus nf4, with calibration methods such as absmax, minmax, rms and fixed range. Quantization is configured with a regex-based rule system that matches Flax module paths, and it integrates with any Flax Linen or NNX model through a single function call, without needing to modify the model's own code.
It is written in Python and installed directly from its GitHub repository.