qwix

What it is and what it does

Qwix is a JAX quantization library that reduces model size and accelerates inference by converting neural network weights and activations to lower-precision numeric types (int4, int8, fp8, and emulated formats). It integrates with Flax models without requiring code changes and supports three main workflows: QAT (fake quantization during training), PTQ (post-training quantization for XLA devices), and ODML (annotation for LiteRT conversion). The library uses a regex-based configuration system to define quantization rules per module, allowing fine-grained control over which layers are quantized and how.

Qwix is designed for practitioners who want to deploy JAX models efficiently on resource-constrained devices or accelerators. It handles the complexity of calibration (absmax, minmax, rms, fixed), granularity (per-channel and sub-channel), and operator-specific quantization strategies for both XLA targets (CPU/GPU/TPU) and mobile/edge targets via LiteRT. The package is actively maintained, in alpha status, and carries no known security vulnerabilities.

Use it for:

• Reduce model size for deployment on mobile or edge devices using LiteRT quantization with full integer arithmetic.
• Speed up inference on TPU/GPU by applying post-training quantization to existing trained models without retraining.
• Train models with quantization awareness using fake quantization to simulate low-precision behavior during training.
• Apply LoRA/QLoRA fine-tuning to quantized models for efficient adaptation to downstream tasks.
• Experiment with different quantization schemas (weight-only, dynamic-range, static-range) via configuration without modifying model code.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Qwix is a JAX quantization library that applies Quantization-Aware Training (QAT) and Post-Training Quantization (PTQ) to neural network models, supporting deployment on XLA devices (CPU/GPU/TPU) and LiteRT targets.

Yes, if you are working with JAX models and need production-ready quantization. The package is actively maintained, has no known vulnerabilities, uses a permissive license, and integrates seamlessly with Flax. Install friction is low for users already in the JAX ecosystem. The alpha status and small community (128 GitHub stars) mean fewer battle-tested recipes, so verify quantization quality for your specific models before production use.

Install

qwix on PyPI

pip install qwix

uv add qwix

poetry add qwix

Installing qwix

Before you install

Low install friction with a pure-Python wheel. The package is actively maintained (last commit 2026-08-13) and in alpha status. It depends on six heavy scientific libraries (jax, jaxlib, flax, numpy, opt_einsum, absl-py), all of which are standard in the JAX ecosystem, so installation complexity is typical for JAX projects rather than exceptional.

License in practice

Licensed under Apache-2.0 (permissive), so you can use, modify, and distribute the package freely in both open-source and commercial projects without copyleft obligations.

Quickstart

pip install qwix

import qwix
import jax
from flax import linen as nn

rules = [qwix.QuantizationRule(module_path='.*', weight_qtype='int8', act_qtype='int8')]
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))

Requires JAX and jaxlib to be installed; jaxlib installation can be complex depending on your hardware (CPU-only, GPU, or TPU). Requires Python 3.10 or later.

Verify before relying

• Whether the package is available on PyPI or still requires installation from GitHub (description says 'doesn't provide a PyPI package yet' but version 0.1.8 appears in PyPI)
• Performance impact of QAT fake quantization versus true quantized training on model convergence
• Compatibility with Flax NNX models beyond Flax Linen

Package facts

Evidence: qwix-0.1.8-py3-none-any.whl

Tags

More Artificial Intelligence packages