### Configure Quantization Rules in Python

Source: https://github.com/google/qwix/blob/main/docs/source/get_started.rst

This code sets up quantization rules for Qwix to apply int8 weight and activation quantization to all modules. It requires the qwix library, takes no runtime inputs, and produces a list of rules. Limited to basic regex matching and does not support dynamic rules or custom quantization types beyond int8.

```python
import qwix

rules = [
    qwix.QuantizationRule(
        module_path='.*',  # this rule matches all modules.
        weight_qtype='int8',  # quantizes weights in int8.
        act_qtype='int8',  # quantizes activations in int8.
    )
]
```

--------------------------------

### Apply Post-Training Quantization in Flax Linen/NNX

Source: https://github.com/google/qwix/blob/main/docs/source/get_started.rst

This code applies quantization to the model using Qwix's PtqProvider. It requires qwix, the model, rules, and for NNX an input sample; outputs a quantized model. Assumes the model is defined and rules are set up; NNX requires a model_input for tracing, while Linen does not.

```python
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))
```

```python
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules), model_input)
```

--------------------------------

### Install Qwix via pip (sh)

Source: https://github.com/google/qwix/blob/main/README.md

Installs the Qwix library directly from its GitHub repository using pip. No PyPI release is available yet, so this command is required before any Python usage. It works on any system with pip installed.

```sh
pip install git+https://github.com/google/qwix
```

--------------------------------

### Qwix Subchannel Quantization for Fine-Grained Control

Source: https://context7.com/google/qwix/llms.txt

Demonstrates subchannel quantization in Qwix, which divides channels into tiles for finer granularity. This is beneficial for layers with high weight variance within channels. The example shows how to configure 'tile_size' for specific layers and contrasts it with per-channel quantization for other layers.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

class LargeModel(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(4096)(x)  # Large layer benefits from subchannel
        x = nn.relu(x)
        x = nn.Dense(4096)(x)
        x = nn.relu(x)
        x = nn.Dense(1000)(x)
        return x

model = LargeModel()
inputs = jax.random.uniform(jax.random.key(0), (16, 512))

# Subchannel quantization with tile_size
subchannel_rules = [
    qwix.QuantizationRule(
        module_path='Dense_0|Dense_1',  # Large layers
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        tile_size=128,  # Divide channels into tiles of 128
    ),
    qwix.QuantizationRule(
        module_path='Dense_2',  # Smaller output layer
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        # No tile_size: uses per-channel quantization
    ),
]

ptq_model = qwix.quantize_model(model, qwix.PtqProvider(subchannel_rules))
fp_params = model.init(jax.random.key(0), inputs)['params']
ptq_abstract = jax.eval_shape(ptq_model.init, jax.random.key(0), inputs)['params']
ptq_params = qwix.quantize_params(fp_params, ptq_abstract)
```

--------------------------------

### Define MLP Model in Flax Linen/NNX

Source: https://github.com/google/qwix/blob/main/docs/source/get_started.rst

This code defines a simple MLP model using Flax's Linen or NNX modules for JAX. It requires jax and flax imports, takes input dimensions and random keys, and outputs model predictions. Assumes basic knowledge of JAX and does not handle advanced features like bias or custom activations.

```python
import jax
from flax import linen as nn

class MLP(nn.Module):

  dhidden: int
  dout: int

  @nn.compact
  def __call__(self, x):
    x = nn.Dense(self.dhidden, use_bias=False)(x)
    x = nn.relu(x)
    x = nn.Dense(self.dout, use_bias=False)(x)
    return x

model = MLP(64, 16)
model_input = jax.random.uniform(jax.random.key(0), (8, 16))
```

```python
import jax
from flax import nnx

class MLP(nnx.Module):

  def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)
    self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)

  def __call__(self, x):
    x = self.linear1(x)
    x = nnx.relu(x)
    x = self.linear2(x)
    return x

model = MLP(16, 64, 16, rngs=nnx.Rngs(0))
model_input = jax.random.uniform(jax.random.key(0), (8, 16))
```

--------------------------------

### Qwix Calibration Methods for Quantization Range Determination

Source: https://context7.com/google/qwix/llms.txt

Illustrates various calibration methods supported by Qwix for computing quantization ranges. These methods include 'absmax', 'minmax', 'minmax' with percentile, 'rms', and 'fixed'. The example also shows how to define different rules for different layer types within a model.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = MLP()

# absmax: Symmetric, uses max absolute value
absmax_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    weight_calibration_method='absmax',
    act_calibration_method='absmax',
)

# minmax: Asymmetric, uses min and max values
minmax_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    weight_calibration_method='minmax',
    act_calibration_method='minmax',
)

# minmax with percentile (clip outliers)
percentile_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    act_calibration_method='minmax,0.99999',
)

# rms: Symmetric, uses root mean square
rms_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    weight_calibration_method='rms,1.0',
    act_calibration_method='rms,2.0',
)

# fixed: User-defined fixed range
fixed_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    act_calibration_method='fixed,-1.0,1.0',
)

# Example: Different calibration per layer type
mixed_calibration_rules = [
    qwix.QuantizationRule(
        module_path=r'Conv.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        weight_calibration_method='absmax',
        act_calibration_method='minmax,0.99999',
    ),
    qwix.QuantizationRule(
        module_path=r'Dense.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        weight_calibration_method='rms,1.0',
        act_calibration_method='absmax',
    ),
]

qt_model = qwix.quantize_model(model, qwix.QtProvider(mixed_calibration_rules))
```

--------------------------------

### Inspect Quantized Parameters in Flax Linen/NNX

Source: https://github.com/google/qwix/blob/main/docs/source/get_started.rst

This console code inspects the quantized model parameters, showing int8 quantized weights with scales. It requires jax or nnx functions and the quantized model; outputs parameter structures. Intended for verification; displays approximate shapes and types, not full values.

```python
>>> jax.eval_shape(ptq_model.init, jax.random.key(0), model_input)['params']
{
  'Dense_0': {
    'kernel': WithAux(
      array=QArray(
        qvalue=ShapeDtypeStruct(shape=(16, 64), dtype=int8),
        scale=ShapeDtypeStruct(shape=(1, 64), dtype=float32),
        ...
      ),
      ...
    )
  },
  'Dense_1': {
    'kernel': WithAux(
      array=QArray(
        qvalue=ShapeDtypeStruct(shape=(64, 16), dtype=int8),
        scale=ShapeDtypeStruct(shape=(1, 16), dtype=float32),
        ...
      ),
      ...
    )
  }
}
```

```python
>>> jax.eval_shape(nnx.to_pure_dict, nnx.state(ptq_model))
{
  'linear1': {
    'kernel': {
      'array': {
        'qvalue': ShapeDtypeStruct(shape=(16, 64), dtype=int8),
        'scale': ShapeDtypeStruct(shape=(1, 64), dtype=float32)
      }
    }
  },
  'linear2': {
    'kernel': {
      'array': {
        'qvalue': ShapeDtypeStruct(shape=(64, 16), dtype=int8),
        'scale': ShapeDtypeStruct(shape=(1, 16), dtype=float32)
      }
    }
  }
}
```

--------------------------------

### Verify Subchannel Structure and FP8 Quantization with Qwix

Source: https://context7.com/google/qwix/llms.txt

This snippet verifies the shape of quantized weights and scales in a subchannel setup, then applies FP8 quantization rules using Qwix's PtqProvider for the entire model. It requires pre-quantized parameters (ptq_params) and a base model, with inputs being regex paths and quantization types. Outputs a quantized model suitable for inference; limitations include dependency on JAX/NumPy and potential accuracy loss without calibration.

```python
dense0_kernel = ptq_params['Dense_0']['kernel'].array
print(f"Weight shape: {dense0_kernel.qvalue.shape}")  # (512, 4096)
print(f"Scale shape: {dense0_kernel.scale.shape}")  # (1, 32) with tile_size=128

# FP8 with subchannel for maximum accuracy
fp8_subchannel_rules = [
    qwix.QuantizationRule(
        module_path='.*',
        weight_qtype=jnp.float8_e4m3fn,
        act_qtype=jnp.float8_e4m3fn,
        tile_size=32,  # Fine-grained quantization
    )
]

fp8_model = qwix.quantize_model(model, qwix.PtqProvider(fp8_subchannel_rules))
```

--------------------------------

### Static-Range Quantization for Linen Models

Source: https://github.com/google/qwix/blob/main/docs/source/ptq.rst

Example of static-range quantization (SRQ) for Linen models using Qwix. This involves calculating static scales from `quant_stats` collected during QAT. It shows how to use `quantize_model` for QAT and PTQ, and `quantize_params` to combine QAT parameters with PTQ scales and stats.

```python
import qwix
import jax
import qconfig

model = SomeLinenModel(...)
rules = [
    qconfig.QuantizationRule(
        weight_qtype="int8",
        act_qtype="int8",
        act_static_scale=True,
    ),
]

qat_model = qwix.quantize_model(model, qwix.QatProvider(rules))
qat_variables = qat_model.init(jax.random.key(0), model_input)
# qat_variables contains "params" and "quant_stats".

ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))
abs_ptq_variables = jax.eval_shape(ptq_model.init, jax.random.key(0), model_input)

ptq_params = qwix.quantize_params(
    qat_variables['params'],
    abs_ptq_variables['params'],
    qat_variables['quant_stats'],
)
```

--------------------------------

### Python - Mobile Deployment with OdmlQatProvider

Source: https://context7.com/google/qwix/llms.txt

Shows how to prepare models for TensorFlow Lite conversion using the OdmlQatProvider. This provider quantizes all operations and adds annotations required for full integer model generation with LiteRT.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

model = MLP()
inputs = jax.random.uniform(jax.random.key(0), (8, 16))

# Configure for ODML with minmax calibration (LiteRT default)
odml_rules = [
    qwix.QuantizationRule(
        module_path='.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        act_static_scale=True,
        act_calibration_method='minmax',
    )
]

# Create ODML QAT model with fixed input range
odml_provider = qwix.OdmlQatProvider(
    odml_rules,
    fixed_range_for_inputs=(-1.0, 1.0),  # Known input range
    disable_per_channel_weights=False,
)

odml_qat_model = qwix.quantize_model(model, odml_provider)

# Train with ODML quantization
variables = odml_qat_model.init(jax.random.key(0), inputs)

# Training loop would collect quant_stats
for epoch in range(10):
    output, new_vars = odml_qat_model.apply(
        variables, inputs, mutable='quant_stats'
    )
    variables = variables.copy({'quant_stats': new_vars['quant_stats']})
```

--------------------------------

### Python - Quantization-Aware Training with QtProvider

Source: https://context7.com/google/qwix/llms.txt

Illustrates quantization-aware training (QAT) using the QtProvider. This approach simulates quantization during training to improve the accuracy of quantized models, incorporating stochastic rounding and statistics collection.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.training import train_state
import optax
import qwix

model = MLP()
inputs = jax.random.uniform(jax.random.key(0), (8, 16))
targets = jax.random.randint(jax.random.key(1), (8,), 0, 10)

# Configure QAT with static-range quantization
qt_rules = [
    qwix.QtRule(
        module_path='.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        act_static_scale=True,  # Collect activation statistics
        bwd_qtype=jnp.int8,  # Quantize gradients
        bwd_stochastic_rounding='zero_mean',  # Reduce quantization error
    )
]

# Create QAT model
qt_provider = qwix.QtProvider(qt_rules)
qt_model = qwix.quantize_model(model, qt_provider)

# Initialize with quant_stats collection
variables = qt_model.init(jax.random.key(0), inputs)

# Training step with mutable quant_stats
def train_step(state, batch):
    def loss_fn(params):
        logits, new_vars = qt_model.apply(
            {'params': params, 'quant_stats': state.quant_stats},
            batch['inputs'],
            mutable='quant_stats'
        )
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits, batch['targets']
        ).mean()
        return loss, new_vars

    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, new_vars), grads = grad_fn(state.params)
    state = state.apply_gradients(grads=grads)
    return state, loss, new_vars['quant_stats']

# After training, convert to PTQ for efficient inference
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(qt_rules))
ptq_abstract = jax.eval_shape(ptq_model.init, jax.random.key(0), inputs)['params']
ptq_params = qwix.quantize_params(
    variables['params'],
    ptq_abstract,
    variables['quant_stats']
)
```

--------------------------------

### Python - Post-Training Quantization with PtqProvider

Source: https://context7.com/google/qwix/llms.txt

Demonstrates post-training quantization (PTQ) using the PtqProvider. This method quantizes weights after model initialization for efficient inference, supporting dynamic-range and static-range quantization.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

# Define model
class ConvNet(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(32, (3, 3))(x)
        x = nn.relu(x)
        x = nn.Conv(64, (3, 3))(x)
        x = nn.relu(x)
        x = x.reshape((x.shape[0], -1))
        x = nn.Dense(10)(x)
        return x

model = ConvNet()
inputs = jax.random.uniform(jax.random.key(0), (4, 28, 28, 1))

# Configure PTQ with per-channel quantization
ptq_rules = [
    qwix.QuantizationRule(
        module_path='Conv.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
    ),
    qwix.QuantizationRule(
        module_path='Dense.*',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
    ),
]

# Create PTQ model
ptq_provider = qwix.PtqProvider(ptq_rules)
ptq_model = qwix.quantize_model(model, ptq_provider)

# Initialize and quantize parameters
fp_params = model.init(jax.random.key(1), inputs)['params']
ptq_abstract = jax.eval_shape(ptq_model.init, jax.random.key(0), inputs)['params']
ptq_params = qwix.quantize_params(fp_params, ptq_abstract)

# Inference
output = ptq_model.apply({'params': ptq_params}, inputs)
print(output.shape)  # (4, 10)
```

--------------------------------

### Apply QLoRA to Transformer model (Python/JAX/Flax)

Source: https://context7.com/google/qwix/llms.txt

Implements parameter-efficient fine-tuning with QLoRA (quantized LoRA) for Transformer models. Shows rule configuration for int4 quantization and LoRA adapter injection. Demonstrates parameter freezing strategy.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

class Transformer(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(256, name='query')(x)
        x = nn.Dense(256, name='key')(x)
        x = nn.Dense(256, name='value')(x)
        x = nn.Dense(128, name='output')(x)
        return x

model = Transformer()
inputs = jax.random.uniform(jax.random.key(0), (4, 64))

# QLoRA: Quantize base model, add LoRA adapters
qlora_rules = [
    qwix.LoraRule(
        module_path='.*',  # Apply to all Dense layers
        weight_qtype=jnp.int4,  # Quantize base weights to int4
        rank=16,  # LoRA rank
        alpha=32.0,  # LoRA alpha (scaling factor = alpha/rank = 2.0)
        dropout=0.1,  # Dropout for LoRA adapters
    )
]

# Create QLoRA model
qlora_provider = qwix.LoraProvider(qlora_rules)
qlora_model = qwix.apply_lora_to_model(model, qlora_provider)

# Initialize parameters
variables = qlora_model.init(jax.random.key(0), inputs)

# Inspect structure: base weights are quantized, LoRA A/B are float
print(variables['params'].keys())
# Output: Contains both quantized base weights and float LoRA adapters

# Training only updates LoRA parameters (freeze base model)
trainable_params = variables['params']['lora']  # Only LoRA weights
frozen_params = variables['params']['base']  # Quantized base weights

# Alternative: LoRA without quantization
lora_rules = [
    qwix.LoraRule(
        module_path='query|key|value',  # Only attention projections
        rank=8,
        alpha=16.0,
    )
]
lora_model = qwix.apply_lora_to_model(model, qwix.LoraProvider(lora_rules))
```

--------------------------------

### Create and manipulate QArray quantized arrays (Python/JAX)

Source: https://context7.com/google/qwix/llms.txt

Demonstrates quantization/dequantization workflows using qwix's QArray. Supports per-channel and subchannel quantization with configurable calibration methods. Shows array operations on quantized data.

```python
import jax
import jax.numpy as jnp
from qwix._src.core import qarray

# Create a floating-point array
fp_array = jax.random.uniform(jax.random.key(0), (64, 128))

# Define quantization strategy
how = qarray.HowToQuantize(
    qtype=jnp.int8,
    channelwise_axes=[1],  # Per-channel along dimension 1
    calibration_method='absmax',
)

# Quantize the array
q_array = qarray.quantize(fp_array, how)
print(f"Quantized shape: {q_array.shape}")  # (64, 128)
print(f"Quantized dtype: {q_array.qvalue.dtype}")  # int8
print(f"Scale shape: {q_array.scale.shape}")  # (1, 128) for per-channel
print(f"Zero point: {q_array.zero_point}")  # None for symmetric

# Dequantize back to floating-point
reconstructed = qarray.dequantize(q_array)
print(f"Reconstruction error: {jnp.abs(fp_array - reconstructed).max()}")

# Subchannel quantization with tiling
subchannel_how = qarray.HowToQuantize(
    qtype=jnp.int8,
    channelwise_axes=[0],
    tiled_axes={1: 32},  # Tile size 32 along dimension 1
    calibration_method='rms,1.0',
)
q_subchannel = qarray.quantize(fp_array, subchannel_how)
print(f"Subchannel scale shape: {q_subchannel.scale.shape}")  # (64, 4) = (64, 128/32)

# Array operations on QArray
reshaped_q = q_array.reshape(128, 64)
transposed_q = q_array.transpose()
sliced_q = q_array[:32, :64]
```

--------------------------------

### Apply post‑training quantization to model (Python)

Source: https://github.com/google/qwix/blob/main/README.md

Runs Qwix's post‑training quantization workflow on the previously defined model using the provided quantization rules. The resulting `ptq_model` contains quantized weights ready for inference on XLA devices.

```python
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))
```

--------------------------------

### Convert model to TFLite using OdmlConversionProvider (Python)

Source: https://context7.com/google/qwix/llms.txt

Quantizes and converts models to TFLite format using qwix's conversion provider. Requires model parameters and quantization statistics as inputs. Outputs a quantized model ready for TFLite conversion.

```python
conversion_provider = qwix.OdmlConversionProvider(
    odml_rules,
    variables['params'],
    variables['quant_stats']
)
conversion_model = qwix.quantize_model(model, conversion_provider)
```

--------------------------------

### FP8 Quantization and Dequantization with Qwix

Source: https://context7.com/google/qwix/llms.txt

Demonstrates how to use Qwix for FP8 quantization of arrays and dequantization of quantized models. It shows the process of quantizing a float array to FP8 and then dequantizing it to measure reconstruction error. It also illustrates extracting and dequantizing specific weights from a quantized model.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

# Assume fp_array is a JAX array
# Assume qarray is imported from qwix

fp8_how = qarray.HowToQuantize(
    qtype=jnp.float8_e4m3fn,
    channelwise_axes=[0],
    calibration_method='absmax',
)
fp8_q = qarray.quantize(fp_array, fp8_how)
fp8_reconstructed = qarray.dequantize(fp8_q)
print(f"FP8 error: {jnp.abs(fp_array - fp8_reconstructed).mean()}")

# Example with a quantized model
class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = MLP()
inputs = jax.random.uniform(jax.random.key(0), (8, 16))
ptq_model = qwix.quantize_model(
    model, qwix.PtqProvider([qwix.QuantizationRule(weight_qtype=jnp.int8)])
)
fp_params = model.init(jax.random.key(0), inputs)['params']
ptq_abstract = jax.eval_shape(ptq_model.init, jax.random.key(0), inputs)['params']
ptq_params = qwix.quantize_params(fp_params, ptq_abstract)

# Extract and dequantize a specific weight
dense0_kernel = ptq_params['Dense_0']['kernel'].array  # This is a QArray
dense0_fp = qwix.dequantize(dense0_kernel)
print(f"Dequantized weight shape: {dense0_fp.shape}")
```

--------------------------------

### Quantize model weights using Qwix (Python)

Source: https://github.com/google/qwix/blob/main/README.md

Demonstrates how to quantize floating‑point parameters separately from the model using Qwix's `quantize_params` function. The quantized parameters are then applied to the PTQ‑quantized model for inference.

```python
# Floating-point params, usually loaded from checkpoints.
fp_params = ...

# Abstract quantized params, which serve as a template for quantize_params.
abs_ptq_params = jax.eval_shape(ptq_model.init, jax.random.key(0), model_input)['params']

# Weight quantization.
ptq_params = qwix.quantize_params(fp abs_ptq_params)

# ptq_params contains the quantized weights and can be consumed by ptq_model.
quantized_model_output = ptq_model.apply({'params': ptq_params}, model_input)
```

--------------------------------

### Qwix QuantizationRule Configuration

Source: https://context7.com/google/qwix/llms.txt

Defines quantization behavior using regex-based pattern matching for modules and operations. Rules specify numeric types, calibration methods, and quantization granularity. This class is used to configure `qwix.quantize_model` and `qwix.quantize_params`.

```python
import jax.numpy as jnp
import qwix

# Weight-only quantization with int4
weight_only_rule = qwix.QuantizationRule(
    module_path='.*Dense.*',  # Match all Dense layers
    weight_qtype=jnp.int4,
)

# Dynamic-range quantization with different types per layer
mixed_precision_rules = [
    qwix.QuantizationRule(
        module_path=r'Dense_0',  # First layer
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
    ),
    qwix.QuantizationRule(
        module_path=r'Dense_\d+',  # Remaining layers
        weight_qtype=jnp.float8_e4m3fn,
        act_qtype=jnp.float8_e4m3fn,
        tile_size=32,  # Subchannel quantization
    ),
]

# Static-range quantization with custom calibration
static_range_rule = qwix.QuantizationRule(
    module_path='.*',
    weight_qtype=jnp.int8,
    act_qtype=jnp.int8,
    act_static_scale=True,
    weight_calibration_method='absmax',
    act_calibration_method='minmax,0.99999',  # 99.999th percentile
    act_batch_axes=(0,),  # Calibrate across batch dimension
)

# Example of applying rules (assuming model and PtqProvider exist)
# model = qwix.quantize_model(model, qwix.PtqProvider(mixed_precision_rules))

```

--------------------------------

### Apply ODML QAT to Model (Linen and NNX)

Source: https://github.com/google/qwix/blob/main/docs/source/odml.rst

Demonstrates quantization-aware training using OdmlQatProvider for two different model front‑ends. The provider inserts FakeQuant ops, respects fusion patterns, and can run in strict mode to ensure all ops are supported.

```python
fp_model = SomeLinenModel(...)
provider = qwix.OdmlQatProvider(rules, strict=True)
qat_model = qwix.quantize_model(fp_model, provider)
# qat_model can be trained as usual.
```

```python
fp_model = SomeNNXModel(...)
provider = qwix.OdmlQatProvider(rules, strict=True)
qat_model = qwix.quantize_model(fp_model, provider, model_input)
# qat_model can be trained as usual.
```

--------------------------------

### Static-Range QAT Rules and Initialization for Linen

Source: https://github.com/google/qwix/blob/main/docs/source/qat.rst

Configures static-range quantization rules for a Linen model and initializes the QAT model to collect quantization statistics. This is crucial for static-range quantization where calibration data is needed.

```python
rules = [
    qwix.QuantizationRule(
        weight_qtype='int8',
        act_qtype='int8',
        act_static_scale=True,
    )
]
qat_model = qwix.quantize_model(model, qwix.QatProvider(rules))
qat_model.init(jax.random.key(0), model_input)['quant_stats']
```

--------------------------------

### Quantize model with Qwix in Python (Linen/NNX)

Source: https://github.com/google/qwix/blob/main/docs/source/odml.rst

Quantizes a floating-point model using Qwix's OdmlConversionProvider. Requires QAT variables for parameters and quantization stats. Outputs a quantized model ready for conversion.

```python
qat_variables = ...  # from QAT.
params = qat_variables['params']
quant_stats = qat_variables['quant_stats']

conversion_provider = qwix.OdmlConversionProvider(rules, params, quant_stats)
conversion_model = qwix.quantize_model(fp_model, conversion_provider)
```

```python
qat_model = ...  # from QAT.
params = nnx.to_pure_dict(nnx.state(qat_model, nnx.Param))
quant_stats = nnx.to_pure_dict(nnx.state(qat_model, qwix.QuantStat)),

conversion_provider = qwix.OdmlConversionProvider(rules, params, quant_stats)
conversion_model = qwix.quantize_model(fp_model, conversion_provider, model_input)
```

--------------------------------

### Mixed-Precision PTQ and QAT with Layer-Specific Rules in Qwix

Source: https://context7.com/google/qwix/llms.txt

Defines a multi-layer Flax model and applies mixed-precision quantization rules for PTQ, verifying dtype per layer after parameter quantization. It extends to QAT with backward precision rules using QtProvider. Depends on JAX, Flax, and Qwix; inputs include model, random keys, and rule lists; outputs quantized models and params. Limitations: Requires initialization for abstract shapes; regex paths must match layer names exactly for targeting.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

class MultiLayerModel(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(512, name='input_proj')(x)
        x = nn.relu(x)
        x = nn.Dense(512, name='hidden1')(x)
        x = nn.relu(x)
        x = nn.Dense(512, name='hidden2')(x)
        x = nn.relu(x)
        x = nn.Dense(10, name='output')(x)
        return x

model = MultiLayerModel()
inputs = jax.random.uniform(jax.random.key(0), (32, 128))

# Mixed-precision strategy
mixed_precision_rules = [
    # Input layer: Keep high precision (fp8)
    qwix.QuantizationRule(
        module_path='input_proj',
        weight_qtype=jnp.float8_e4m3fn,
        act_qtype=jnp.float8_e4m3fn,
    ),
    # Hidden layers: Aggressive quantization (int4)
    qwix.QuantizationRule(
        module_path='hidden1|hidden2',
        weight_qtype=jnp.int4,
        act_qtype=jnp.int8,  # Activations more sensitive, use int8
        tile_size=64,  # Subchannel to maintain accuracy with int4
    ),
    # Output layer: High precision for final predictions
    qwix.QuantizationRule(
        module_path='output',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
    ),
]

ptq_model = qwix.quantize_model(model, qwix.PtqProvider(mixed_precision_rules))

# Verify different dtypes per layer
fp_params = model.init(jax.random.key(0), inputs)['params']
ptq_abstract = jax.eval_shape(ptq_model.init, jax.random.key(0), inputs)['params']
ptq_params = qwix.quantize_params(fp_params, ptq_abstract)

print("Input projection:", ptq_params['input_proj']['kernel'].array.qvalue.dtype)
print("Hidden layer 1:", ptq_params['hidden1']['kernel'].array.qvalue.dtype)
print("Output:", ptq_params['output']['kernel'].array.qvalue.dtype)

# Mixed QAT with different backward precision
mixed_qat_rules = [
    qwix.QtRule(
        module_path='input_proj',
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
        bwd_qtype=jnp.float8_e5m2,  # Higher precision for gradients
    ),
    qwix.QtRule(
        module_path='hidden.*',
        weight_qtype=jnp.int4,
        act_qtype=jnp.int8,
        bwd_qtype=jnp.int8,
        tile_size=32,
    ),
]

qt_model = qwix.quantize_model(model, qwix.QtProvider(mixed_qat_rules))
```

--------------------------------

### Quantize Linen Model with PTQ

Source: https://github.com/google/qwix/blob/main/docs/source/ptq.rst

Quantizes a Linen model using Qwix's PtqProvider. Requires a floating-point model and quantization rules as input. Outputs a quantized model ready for deployment.

```python
fp_model = SomeLinenModel()
ptq_model = qwix.quantize_model(fp_model, qwix.PtqProvider(rules))
```

--------------------------------

### Static-Range QAT Rules and Access for NNX

Source: https://github.com/google/qwix/blob/main/docs/source/qat.rst

Configures static-range quantization rules for an NNX model and applies QAT. It then demonstrates accessing the collected quantization statistics from a specific layer's weights.

```python
rules = [
    qwix.QuantizationRule(
        weight_qtype='int8',
        act_qtype='int8',
        act_static_scale=True,
    )
]
qat_model = qwix.quantize_model(model, qwix.QatProvider(rules), model_input)
qat_model.linear1.dot_general0_lhs
```

--------------------------------

### Quantize Parameters for PTQ with Qwix

Source: https://context7.com/google/qwix/llms.txt

Converts floating-point parameters to a quantized format for Post-Training Quantization (PTQ). This function operates on parameter trees and requires an abstract quantized parameter structure as a template. It takes floating-point parameters, the abstract quantized structure, and optionally calibration statistics as input, returning quantized parameters.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

# Setup model and quantization
model = MLP() # Assuming MLP is defined as in the previous example
inputs = jax.random.uniform(jax.random.key(0), (8, 16))
rules = [qwix.QuantizationRule(weight_qtype=jnp.int8, act_qtype=jnp.int8)]
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))

# Initialize floating-point parameters (e.g., from checkpoint)
fp_params = model.init(jax.random.key(42), inputs)['params']

# Get abstract quantized parameter structure
ptq_abstract_params = jax.eval_shape(
    ptq_model.init, jax.random.key(0), inputs
)['params']

# Quantize the parameters
ptq_params = qwix.quantize_params(fp_params, ptq_abstract_params)

# Run inference with quantized parameters
output = ptq_model.apply({'params': ptq_params}, inputs)
print(output.shape)  # (8, 10)

# For static-range quantization, include quant_stats
# ptq_params = qwix.quantize_params(fp_params, ptq_abstract_params, quant_stats)

```

--------------------------------

### Quantize Flax Model with Qwix

Source: https://context7.com/google/qwix/llms.txt

Transforms a Flax model into a quantized version using specified providers and rules. This is the primary interface for applying quantization to Flax Linen or NNX models without code modification. It takes a model instance and a provider (e.g., PtqProvider) as input, returning a quantized model.

```python
import jax
import jax.numpy as jnp
from flax import linen as nn
import qwix

# Define a simple model
class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(64, use_bias=False)(x)
        x = nn.relu(x)
        x = nn.Dense(10, use_bias=False)(x)
        return x

model = MLP()
inputs = jax.random.uniform(jax.random.key(0), (8, 16))

# Configure quantization rules
rules = [
    qwix.QuantizationRule(
        module_path='.*',  # Match all modules
        weight_qtype=jnp.int8,
        act_qtype=jnp.int8,
    )
]

# Create quantized model
ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules))

# Verify quantized structure
abstract_params = jax.eval_shape(
    ptq_model.init, jax.random.key(0), inputs
)['params']
print(abstract_params)
# Output: Parameters wrapped in WithAux containing QArray objects

```

--------------------------------

### Define ODML Quantization Rules

Source: https://github.com/google/qwix/blob/main/docs/source/odml.rst

Creates a list of QuantizationRule objects specifying int8 weight and activation types, which is used by ODML providers for static-range quantization. No external dependencies beyond the qwix library are required.

```python
rules = [
    qwix.QuantizationRule(
        weight_qtype='int8',
        act_qtype='int8',
    )
]
```

--------------------------------

### Convert model to LiteRT format in Python (Linen/NNX)

Source: https://github.com/google/qwix/blob/main/docs/source/odml.rst

Converts a quantized model to LiteRT format using ai_edge_jax. Requires model parameters and input. Outputs a LiteRT model that can be evaluated and exported.

```python
import ai_edge_jax

litert_model = ai_edge_jax.convert(
    conversion_model.apply,
    {'params': params},
    (model_input,),
    _litert_converter_flags={'_experimental_strict_qdq': True},  # necessary for Qwix.
)

# Evaluate the LiteRT model on the host.
litert_result = litert_model(model_input)
# Export the LiteRT model.
litert_model.export('/tmp/litert_model.tflite')
```

```python
import ai_edge_jax

graphdef, state = nnx.split(conversion_model)
litert_model = ai_edge_jax.convert(
    lambda params, *args: nnx.merge(graphdef, params)(*args),
    state,
    (model_input,),
    _litert_converter_flags={'_experimental_strict_qdq': True},  # necessary for Qwix.
)

# Evaluate the LiteRT model on the host.
litert_result = litert_model(model_input)
# Export the LiteRT model.
litert_model.export('/tmp/litert_model.tflite')
```

--------------------------------

### Quantize NNX Model with PTQ

Source: https://github.com/google/qwix/blob/main/docs/source/ptq.rst

Quantizes an NNX model using Qwix's PtqProvider. Requires a floating-point model, quantization rules, and model input. Outputs a quantized model. JIT can be used to eliminate intermediate floating-point model.

```python
fp_model = SomeNnxModel()
ptq_model = qwix.quantize_model(fp_model, qwix.PtqProvider(rules), model_input)
```

```python
def create_quantized_model():
  fp_model = SomeNnxModel()
  return qwix.quantize_model(fp_model, qwix.PtqProvider(rules), model_input)

ptq_model = nnx.jit(create_quantized_model)()
```

--------------------------------

### Static-Range Quantization for NNX Models

Source: https://github.com/google/qwix/blob/main/docs/source/ptq.rst

This snippet illustrates static-range quantization (SRQ) for NNX models with Qwix. It demonstrates quantizing a QAT model to PTQ, where `quantize_model` handles the conversion of `quant_stats`. It also shows using `quantize_params` for NNX models, extracting state and parameters from QAT and PTQ models.

```python
import qwix
import nnx
import qconfig

model = SomeNnxModel(...)
rules = [
    qconfig.QuantizationRule(
        weight_qtype="int8",
        act_qtype="int8",
        act_static_scale=True,
    ),
]

qat_model = qwix.quantize_model(model, qwix.QatProvider(rules), model_input)
# qat_model contains both params and quant_stats.

# quantize_model converts the quant_stats if the PTQ model is converted from
# a QAT model.
ptq_model = qwix.quantize_model(qat_model, qwix.PtqProvider(rules), model_input)

# It's also possible to use quantize_params for NNX models.
ptq_params = qwix.quantize_params(
    nnx.to_pure_dict(nnx.state(qat_model, nnx.Param)),
    ptq_model,  # or abs_ptq_model
    nnx.to_pure_dict(nnx.state(qat_model, qwix.QuantStat)),
)
```

--------------------------------

### Enable int8 dynamic-range quantization in Qwix

Source: https://github.com/google/qwix/blob/main/docs/source/basics.rst

Configures Qwix to quantize both weights and activations to int8. Activations are quantized dynamically during inference using optimal scales calculated from the data.

```python
qwix.QuantizationRule(
    weight_qtype="int8",
    act_qtype="int8",
)
```

--------------------------------

### Apply LoRA/QLoRA with NNX

Source: https://github.com/google/qwix/blob/main/docs/source/lora.rst

Applies LoRA/QLoRA to a model using Qwix's LoraProvider within the NNX framework. It defines LoraRule with weight quantization type, rank, and alpha, then applies it to the model. The output demonstrates the structure of the model's state after applying LoRA.

```python
rules = [
    qwix.LoraRule(
        weight_qtype='nf4',
        rank=16,
        alpha=0.5,
    )
]
lora_model = qwix.apply_lora_to_model(model, qwix.LoraProvider(rules), model_input)
```

```py
>>> jax.eval_shape(nnx.to_pure_dict, nnx.state(lora_model))
{'linear1': {'kernel': {'array': {'qvalue': ShapeDtypeStruct(shape=(16, 64), dtype=uint4),
   'scale': ShapeDtypeStruct(shape=(1, 64), dtype=float32)}}},
 'kernel_lora_a': ShapeDtypeStruct(shape=(16, 16), dtype=float32),
 'kernel_lora_b': ShapeDtypeStruct(shape=(16, 64), dtype=float32)},
 'linear2': {'kernel': {'array': {'qvalue': ShapeDtypeStruct(shape=(64, 16), dtype=uint4),
   'scale': ShapeDtypeStruct(shape=(1, 16), dtype=float32)}}},
 'kernel_lora_a': ShapeDtypeStruct(shape=(64, 16), dtype=float32),
 'kernel_lora_b': ShapeDtypeStruct(shape=(16, 16), dtype=float32)}
```

--------------------------------

### Apply QAT to Linen Model

Source: https://github.com/google/qwix/blob/main/docs/source/qat.rst

Applies Quantization-Aware Training (QAT) to a Linen model using Qwix's quantize_model function and QatProvider. This process emulates quantized operations using floating-point math.

```python
fp_model = SomeLinenModel(...)
qat_model = qwix.quantize_model(fp_model, qwix.QatProvider(rules))
```

--------------------------------

### Dequantize QArray to floating-point (Python/JAX)

Source: https://context7.com/google/qwix/llms.txt

Reconstructs original floating-point arrays from quantized QArray format. Compares symmetric (absmax) and asymmetric (minmax) quantization methods. Shows error measurement for reconstruction quality.

```python
import jax
import jax.numpy as jnp
from qwix._src.core import qarray
import qwix

# Create and quantize an array
fp_array = jax.random.normal(jax.random.key(0), (32, 64))

# Symmetric quantization (absmax)
symmetric_how = qarray.HowToQuantize(
    qtype=jnp.int8,
    channelwise_axes=[],  # Per-tensor
    calibration_method='absmax',
)
symmetric_q = qarray.quantize(fp_array, symmetric_how)
symmetric_reconstructed = qarray.dequantize(symmetric_q)
print(f"Symmetric error: {jnp.abs(fp_array - symmetric_reconstructed).mean()}")

# Asymmetric quantization (minmax)
asymmetric_how = qarray.HowToQuantize(
    qtype=jnp.int8,
    channelwise_axes=[1],  # Per-channel
    calibration_method='minmax',
)
asymmetric_q = qarray.quantize(fp_array, asymmetric_how)
print(f"Has zero_point: {asymmetric_q.zero_point is not None}")  # True
asymmetric_reconstructed = qarray.dequantize(asymmetric_q)
```

--------------------------------

### Define MLP model with Flax Linen (Python)

Source: https://github.com/google/qwix/blob/main/README.md

Creates a simple multilayer perceptron using Flax's Linen API. The model is pure‑functional and can be quantized without modifications. Returns a callable model and a random input tensor for demonstration.

```python
import jax
from flax import linen as nn

class MLP(nn.Module):

  dhidden: int
  dout: int

  @nn.compact
  def __call__(self, x):
    x = nn.Dense(self.dhidden, use_bias=False)(x)
    x = nn.relu(x)
    x = nn.Dense(self.dout, use_bias=False)(x)
    return x

model = MLP(64, 16)
model_input = jax.random.uniform(jax.random.key(0), (8, 16))
```