{"id":6219,"library":"qwix","title":"Qwix: A JAX Quantization Library","description":"Qwix is a JAX-native quantization library for both research and production, providing efficient model compression through Quantization-Aware Training (QAT), Post-Training Quantization (PTQ), and ODML quantization. It supports various XLA targets (CPU/GPU/TPU) and LiteRT, featuring a flexible, regex-based configuration system for Flax Linen and NNX models. The current version is 0.1.5, with an active release cadence.","status":"active","version":"0.1.5","language":"en","source_language":"en","source_url":"https://github.com/google/qwix","tags":["jax","quantization","machine-learning","deep-learning","flax","nnx"],"install":[{"cmd":"pip install git+https://github.com/google/qwix","lang":"bash","label":"Install from GitHub (recommended by project)"}],"dependencies":[{"reason":"Core dependency for JAX-native operations.","package":"jax","optional":false},{"reason":"Used for model definition and integration with Flax Linen/NNX models.","package":"flax","optional":false},{"reason":"Requires Python 3.10 or higher.","package":"python","optional":false}],"imports":[{"symbol":"qwix","correct":"import qwix"},{"note":"QuantizationRule is typically accessed as `qwix.QuantizationRule` after `import qwix`.","wrong":"import qwix.QuantizationRule","symbol":"QuantizationRule","correct":"from qwix import QuantizationRule"},{"note":"quantize_model is typically accessed as `qwix.quantize_model` after `import qwix`.","wrong":"import qwix.quantize_model","symbol":"quantize_model","correct":"from qwix import quantize_model"}],"quickstart":{"code":"import jax\nimport jax.numpy as jnp\nfrom flax import linen as nn\nimport qwix\n\n# Define a simple MLP model using Flax Linen\nclass MLP(nn.Module):\n    dhidden: int = 64\n    dout: int = 10\n\n    @nn.compact\n    def __call__(self, x):\n        x = nn.Dense(self.dhidden, use_bias=False)(x)\n        x = nn.relu(x)\n        x = nn.Dense(self.dout, use_bias=False)(x)\n        return x\n\n# Initialize the model and dummy input\nmodel = MLP()\nkey = jax.random.key(0)\nmodel_input = jax.random.uniform(key, (8, 16))\nparams = model.init(key, model_input)['params']\n\n# Define quantization rules for int8 weight and activation quantization\n# This rule matches all modules ('.*')\nrules = [\n    qwix.QuantizationRule(\n        module_path='.*',\n        weight_qtype=jnp.int8,\n        act_qtype=jnp.int8,\n    )\n]\n\n# Apply Post-Training Quantization (PTQ)\nptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))\n\nprint(\"Original model parameters (example kernel shape):\", params['Dense_0']['kernel'].shape)\nprint(\"Quantized model parameters (example kernel):\", jax.eval_shape(ptq_model.apply, {'params': params}, model_input)['Dense_0']['kernel'])","lang":"python","description":"This quickstart demonstrates how to define a simple Flax MLP model and then apply Post-Training Quantization (PTQ) using Qwix. It configures a `QuantizationRule` to quantize both weights and activations to int8 across all modules in the model, and then applies the quantization using `qwix.quantize_model`."},"warnings":[{"fix":"Always install `qwix` directly from its GitHub repository using `pip install git+https://github.com/google/qwix`.","message":"The PyPI project page for Qwix currently states, 'Qwix doesn't provide a PyPI package yet. To use Qwix, you need to install from GitHub directly.' This is despite the `qwix` package being available on PyPI. For consistent and recommended installation, use the `pip install git+https://github.com/google/qwix` command.","severity":"gotcha","affected_versions":"All versions"},{"fix":"Upgrade to Qwix version 0.1.5 or newer to resolve the LoRA RNG issue.","message":"An RNG issue for LoRA (Low-Rank Adaptation) was fixed in version 0.1.5. If you are using LoRA or QLoRA with Qwix in older versions, you may encounter unexpected behavior related to random number generation.","severity":"gotcha","affected_versions":"<0.1.5"},{"fix":"Currently, there is no direct programmatic way to retrieve the Qwix version from the installed package. You must rely on installation logs or environment information.","message":"Qwix does not currently expose a `__version__` attribute, which is a common Python practice for programmatic version checking. This may complicate dependency management or conditional logic based on the installed Qwix version.","severity":"gotcha","affected_versions":"All versions"},{"fix":"Carefully review the Qwix documentation on 'Training with Quantization (QAT/QT)' to select the appropriate training method for your use case.","message":"Users should understand the distinction between Qwix's Quantization-Aware Training (QAT) and Quantized Training (QT). QAT uses fake quantization to recover model quality by making the model aware of precision loss during inference, while QT performs computations using low-precision integer arithmetic in both forward and backward passes for performance benefits. Choosing the correct mode depends on your specific optimization goals.","severity":"gotcha","affected_versions":"All versions"}],"env_vars":null,"last_verified":"2026-04-14T00:00:00.000Z","next_check":"2026-07-13T00:00:00.000Z","problems":[]}