prepare_fx(model: nn.Module, qconfig_mapping: QConfig | QConfigMapping | None = None, example_inputs: Tensor | None = None, inplace: bool = False)Graph-mode prepare — auto-wrap the quant boundaries and insert observers.
prepare_fx is the graph-mode entry point of static PTQ, the counterpart
to eager lucid.quantization.prepare. It removes the one piece of
manual bookkeeping the eager flow demands: instead of hand-placing a
QuantStub at the model input and a DeQuantStub at the output, it wraps
the model in a lucid.nn.quantized.QuantWrapper so the input is
quantized and the output dequantized automatically, then runs the ordinary
eager lucid.quantization.prepare underneath to attach activation
observers. From there the workflow is identical — run calibration data
through the returned model, then convert_fx.
The graph-mode payoff is downstream: because every quantized replacement is
an ordinary composite, the converted model traces cleanly through
lucid.compile, letting the dispatch-bound quantized graph fuse into a
single executable — the regime where quantization's inference win is
structurally permanent. (True FX-style observer insertion into a traced
graph would need C++ TraceGraph mutation; quantizing at the module level and
letting compile capture it reaches the same result.)
Parameters
modelnn.ModuleQuantStub / DeQuantStub needed;
the wrapper adds the boundaries.lucid.quantization.prepare; None (default) uses the
static-PTQ default mapping.inplacebool= FalseTrue mutate and return model itself; if False (default)
prepare a deep copy.Returns
nn.ModuleThe prepared, eval-mode model (input/output boundaries wrapped,
observers attached). Calibrate it, then call convert_fx.
Notes
- The only behavioral difference from eager
lucid.quantization.prepareis the automaticQuantWrapper— everything else (observer types, calibration semantics, deep-copy default) is shared. example_inputsis a placeholder today; passing it changes nothing.- Nothing is quantized yet — calibrate before
convert_fx, exactly as in the eager flow.
Examples
>>> import lucid, lucid.nn as nn
>>> import lucid.quantization as Q
>>> model = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10))
>>> model.qconfig = Q.get_default_qconfig()
>>> prepared = Q.prepare_fx(model) # no QuantStub placed by hand
>>> for _ in range(8): # calibrate
... _ = prepared(lucid.randn(32, 128))
>>> qmodel = Q.convert_fx(prepared)
>>> compiled = lucid.compile(qmodel) # quantized graph fuses cleanlySee Also
- lucid.quantization.convert_fx—The graph-mode bake step.
- lucid.quantization.prepare—The eager-mode counterpart (manual boundaries).
- lucid.nn.quantized.QuantWrapper—The input/output boundary wrapper used here.