qwix

import jax import jax.numpy as jnp from flax import linen as nn import qwix # Define a simple MLP model using Flax Linen class MLP(nn.Module): dhidden: int = 64 dout: int = 10 @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 # Initialize the model and dummy input model = MLP() key = jax.random.key(0) model_input = jax.random.uniform(key, (8, 16)) params = model.init(key, model_input)['params'] # Define quantization rules for int8 weight and activation quantization # This rule matches all modules ('.*') rules = [ qwix.QuantizationRule( module_path='.*', weight_qtype=jnp.int8, act_qtype=jnp.int8, ) ] # Apply Post-Training Quantization (PTQ) ptq_model = qwix.quantize_model(model, qwix.PtqProvider(rules)) print("Original model parameters (example kernel shape):", params['Dense_0']['kernel'].shape) print("Quantized model parameters (example kernel):", jax.eval_shape(ptq_model.apply, {'params': params}, model_input)['Dense_0']['kernel'])