How does Qwix quantize a model without changing its code?

By intercepting the operations the model calls, not the model. When you call qwix.quantize_model, Qwix installs an interceptor that monkey-patches the JAX and Flax functions a model uses to do its heavy math — jax.lax.dot_general, jax.numpy.einsum, jax.lax.conv_general_dilated, flax.linen.Module.param, and a few more. For the duration of one forward pass, calling dot_general runs Qwix's quantized version instead of the real one. The model's source line still reads nn.Dense(64)(x); it just resolves to a different function underneath. The README puts it plainly: Qwix "integrates with models without need to modify their code, so any model can be used."

What is the difference between PTQ and QAT in Qwix?

They are two providers riding the same interception mechanism. Both patch the same set of ops; they differ only in what the patched op does. PtqProvider (Post-Training Quantization) actually stores weights as low-bit QArrays and runs the matmul in integer arithmetic, then dequantizes — that is the fast serving path on TPU/GPU. QtProvider (the QAT path) does "fake quantization": it quantizes and immediately dequantizes inside the forward pass so the network feels the rounding error during training but the math stays in floating point. Same rails, different car — you swap PtqProvider(rules) for a QAT provider and nothing else about the wiring changes.

How does Qwix decide which layers to quantize?

With a list of QuantizationRule objects, each a regex plus a config. The regex (module_path, default ".*") is matched against every module's path as the model runs; the first rule that matches supplies the quantization types (weight_qtype, act_qtype), the calibration method, and the tile size. A rule with module_path=".*" quantizes everything; a more specific regex targets one block. Qwix deliberately ships no preset recipes: "different quantization schemas are achieved by combinations of quantization configs." get_unused_rules() afterwards tells you which rules never matched, so a typo'd regex is caught rather than silently ignored.

Does Qwix quantize weights, activations, or both?

It depends on the rule, and Qwix decides per operand as each op fires. Inside a patched dot_general, the operand that is a Flax parameter (detected by find_param) is the weight — it is quantized whenever weight_qtype is set. The other operand is an activation, and it is only quantized if act_qtype is also set. So "weight-only" quantization (int4 weights, fp16 activations) is just a rule with act_qtype left None. Granularity is per-op: dot_general and einsum support per-channel and sub-channel scales; conv_general_dilated is per-channel.

What are int8, int4, fp8, and nf4 in Qwix and how do they differ?

They are the numeric grids a weight can be rounded onto. int8 and int4 are uniform integer grids — evenly spaced steps, with symmetric bounds of 127.5 and 7.5 in Qwix's calibration. fp8 (float8_e4m3fn) is a floating grid with a huge dynamic range (max 448) but only three mantissa bits, so its steps get coarser as magnitude grows. nf4 (NormalFloat-4) is non-uniform: sixteen fixed buckets, quantiles of a normal distribution, clustered densely near zero — which is exactly where a trained weight row concentrates, so nf4 beats plain int4 on Gaussian weights. Because nf4 is a lookup table, Qwix notes it "use[s] a non-linear lookup table of buckets, making direct arithmetic on quantized indices mathematically invalid" — so it cannot run the fast integer-matmul path.

How Qwix Quantizes Any Flax Model Without Touching Its Code

You have a trained Flax model and you want it in int8 — but the model code isn't yours to edit. Qwix quantizes it anyway, by a trick that has nothing to do with the model: it swaps the ops the model calls. For the duration of one forward pass it patches jax.lax.dot_general and jnp.einsum, so every matmul routes through a quantized version while the model's own source never changes. We drive the op-swap live, compute int8/int4/fp8/nf4 grids from Qwix's real bounds, show how PTQ and QAT ride the exact same mechanism, and map what gets quantized where — plus the honest costs of patching a running program.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of How Qwix Quantizes Any Flax Model Without Touching Its Code — built to be understood, not skimmed.

Questions

How does Qwix quantize a model without changing its code?: By intercepting the operations the model calls, not the model. When you call qwix.quantize_model, Qwix installs an interceptor that monkey-patches the JAX and Flax functions a model uses to do its heavy math — jax.lax.dot_general, jax.numpy.einsum, jax.lax.conv_general_dilated, flax.linen.Module.param, and a few more. For the duration of one forward pass, calling dot_general runs Qwix's quantized version instead of the real one. The model's source line still reads nn.Dense(64)(x); it just resolves to a different function underneath. The README puts it plainly: Qwix "integrates with models without need to modify their code, so any model can be used."
What is the difference between PTQ and QAT in Qwix?: They are two providers riding the same interception mechanism. Both patch the same set of ops; they differ only in what the patched op does. PtqProvider (Post-Training Quantization) actually stores weights as low-bit QArrays and runs the matmul in integer arithmetic, then dequantizes — that is the fast serving path on TPU/GPU. QtProvider (the QAT path) does "fake quantization": it quantizes and immediately dequantizes inside the forward pass so the network feels the rounding error during training but the math stays in floating point. Same rails, different car — you swap PtqProvider(rules) for a QAT provider and nothing else about the wiring changes.
How does Qwix decide which layers to quantize?: With a list of QuantizationRule objects, each a regex plus a config. The regex (module_path, default ".*") is matched against every module's path as the model runs; the first rule that matches supplies the quantization types (weight_qtype, act_qtype), the calibration method, and the tile size. A rule with module_path=".*" quantizes everything; a more specific regex targets one block. Qwix deliberately ships no preset recipes: "different quantization schemas are achieved by combinations of quantization configs." get_unused_rules() afterwards tells you which rules never matched, so a typo'd regex is caught rather than silently ignored.
Does Qwix quantize weights, activations, or both?: It depends on the rule, and Qwix decides per operand as each op fires. Inside a patched dot_general, the operand that is a Flax parameter (detected by find_param) is the weight — it is quantized whenever weight_qtype is set. The other operand is an activation, and it is only quantized if act_qtype is also set. So "weight-only" quantization (int4 weights, fp16 activations) is just a rule with act_qtype left None. Granularity is per-op: dot_general and einsum support per-channel and sub-channel scales; conv_general_dilated is per-channel.
What are int8, int4, fp8, and nf4 in Qwix and how do they differ?: They are the numeric grids a weight can be rounded onto. int8 and int4 are uniform integer grids — evenly spaced steps, with symmetric bounds of 127.5 and 7.5 in Qwix's calibration. fp8 (float8_e4m3fn) is a floating grid with a huge dynamic range (max 448) but only three mantissa bits, so its steps get coarser as magnitude grows. nf4 (NormalFloat-4) is non-uniform: sixteen fixed buckets, quantiles of a normal distribution, clustered densely near zero — which is exactly where a trained weight row concentrates, so nf4 beats plain int4 on Gaussian weights. Because nf4 is a lookup table, Qwix notes it "use[s] a non-linear lookup table of buckets, making direct arithmetic on quantized indices mathematically invalid" — so it cannot run the fast integer-matmul path.

How Qwix Quantizes Any Flax Model Without Touching Its Code

Questions

Related explainers