fake_quantize(x: Tensor, scale: _ScaleLike, zero_point: _ZeroPointLike, quant_min: int, quant_max: int, ch_axis: int | None = None)Simulate quantization in the float domain with a straight-through gradient.
fake_quantize composes quantize and dequantize into a
single float→float round-trip: it quantizes x to the integer grid and
immediately decodes it back, returning a float32 tensor that carries the
rounding error quantization would introduce while staying differentiable.
That is what makes quantization-aware training (QAT) possible — a network can
feel the low-precision grid on every forward, and thus learn weights that
are robust to it, while gradients continue to flow and the optimizer keeps
working in float. It is the operation behind FakeQuantize, every
nn.qat layer, and the output re-quantization step of the sidecar
quantized inference modules.
The forward is the encode→decode round-trip:
The round has no useful derivative, so the backward is the straight-through estimator (STE): the gradient passes unchanged where the pre-clip code lands inside the grid and is zeroed where the value saturates:
Because of that custom backward it is implemented as a
lucid.autograd.Function rather than a plain composite; only x
is registered as a differentiable input (the qparams / bounds are passed as
keyword arguments), so a single gradient edge flows back to x alone.
Parameters
scaleTensor or floats (typically produced by an observer / learnable
in QAT). A length-C tensor for per-channel.zero_pointTensor, float, or intz.quant_minint0, 255 for quint8, -128, 127 for qint8); values rounding outside this range saturate.quant_maxint0, 255 for quint8, -128, 127 for qint8); values rounding outside this range saturate.ch_axisint= NoneNone (default) uses a
single scalar scale / zero_point for the whole tensor.Returns
TensorA fake-quantized float32 tensor, the same shape as x, carrying
the STE gradient described above.
Notes
- Output stays
float32— nothing is physically stored as int here. Usequantizewhen you actually want integer codes; use this when you want the effect of quantization inside a differentiable graph. - The STE saturation mask is computed from the pre-clip code, so it correctly zeroes the gradient for inputs that round to exactly the grid edge from outside.
- At inference the round-trip makes a module's output bit-equivalent to a
true integer kernel without needing one — the reason the sidecar
lucid.nn.quantized.Linearfake-quantizes its output. - Runs on the public
lucid.*op surface (H4), so the STE is identical on the Accelerate and MLX streams.
Examples
>>> import lucid
>>> import lucid.quantization as Q
>>> x = lucid.tensor([0.02, 0.31, 0.77], requires_grad=True)
>>> y = Q.fake_quantize(x, scale=0.1, zero_point=0, quant_min=0, quant_max=255)
>>> y.shape
(3,)
STE in action — the gradient is 1 in-range and 0 where the input
saturates past quant_max:
>>> x = lucid.tensor([0.5, 99.0], requires_grad=True) # 99.0 rounds past 255*0.1
>>> Q.fake_quantize(x, scale=0.1, zero_point=0, quant_min=0, quant_max=255).sum().backward()
>>> [float(g) for g in x.grad]
[1.0, 0.0]See Also
- lucid.quantization.quantize—The (non-differentiable) integer encode.
- lucid.quantization.dequantize—The affine decode half of the round-trip.
- lucid.quantization.FakeQuantize—Observer-driven module wrapper of this op.
- lucid.quantization.prepare_qat—Inserts fake-quant layers for QAT.