aqtp

import jax import jax.numpy as jnp from flax import linen as nn import aqt.jax.v2 as aqt # Define a simple MLP class Mlp(nn.Module): num_layers: int features: int @nn.compact def __call__(self, x): for i in range(self.num_layers - 1): x = nn.Dense(self.features)(x) x = nn.relu(x) x = nn.Dense(self.features)(x) return x # Create an AQT configuration for int8 quantization aqt_config = aqt.config.DotGeneral( rhs=aqt.config.Tensor(bits=8), lhs=aqt.config.Tensor(bits=8), ) # Function to quantize a module def quantize_module(module, aqt_cfg): return aqt.config.set_dot_general_by_config(module, aqt_cfg) # Initialize model and quantize key = jax.random.PRNGKey(0) model = Mlp(num_layers=3, features=10) input_shape = (1, 5) # Batch size 1, 5 input features params = model.init(key, jnp.zeros(input_shape)) quantized_model = quantize_module(model, aqt_config) quantized_params = quantized_model.init(key, jnp.zeros(input_shape)) print(f"Original output shape: {model.apply(params, jnp.zeros(input_shape)).shape}") print(f"Quantized output shape: {quantized_model.apply(quantized_params, jnp.zeros(input_shape)).shape}")