How to use aqt int8 quantization with flax.nnx

[lkwq007]

I was wondering if there are any examples demonstrating how to use aqt with flax.nnx, specifically for quantizing a pretrained model into a quantized version. Using nnx.bridge.ToLinen to wrap the model and then doing the quantization seems to be a feasible approach, but is there a better way to achieve this? Thanks.

[cgarciae]

Hey @lkwq007, we still haven't ported any of the quantization utilities. Quantization APIs might be different in NNX to leverage the the ability to do model surgery. I'd imagine having a proper LinearFP8 that you can use to monkey patch Linear layers in an existing model.