Quantization

Reducing model precision (float32 → int8 / float16 / bfloat16 / int4) to make inference smaller and faster. TF has multiple quantization stacks; the recommended path depends on the target.

What quantization paths exist

Path	Code	Target
TFLite post-training quantization	tensorflow/lite/tools/optimize/, tensorflow/lite/python/	Mobile .tflite deployment.
TFLite quantization-aware training (QAT)	tensorflow/python/keras/ (with TF Model Optimization toolkit)	Mobile .tflite deployment with higher accuracy.
MLIR-based quantization (tf_quant dialect)	tensorflow/compiler/mlir/quantization/	TFLite + StableHLO + server-side quantized inference.
Server-side INT8 via TensorRT	tensorflow/compiler/tf2tensorrt/	NVIDIA GPU inference with INT8 calibration.
Mixed precision (bfloat16/float16)	tensorflow/python/keras/mixed_precision/, Grappler auto_mixed_precision	Training with reduced-precision compute.
Auto-cast for TPU	TPU bfloat16 paths in tensorflow/python/tpu/	TPU training in bfloat16.

This page describes how the pieces relate; for code-level entry points see the linked subdirectories.

Post-training quantization for TFLite

converter = tf.lite.TFLiteConverter.from_saved_model("path")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = my_dataset_gen   # for full integer
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
tflite_model = converter.convert()

What happens under the hood:

1. The Python API loads the SavedModel and runs the MLIR-based converter (tensorflow/compiler/mlir/lite/).
2. With a representative dataset, the calibration pass observes activation ranges (tensorflow/compiler/mlir/lite/quantization/).
3. Quantization passes rewrite the MLIR tfl graph to use INT8 ops with the discovered scale/zero-point parameters.
4. The flatbuffer is emitted with quantized tensors.

The older TOCO-based path (tensorflow/lite/toco/) is still around but no longer the default.

Quantization-aware training

QAT adds fake-quant ops in the model during training so it learns to be robust to quantization noise. The Python API lives in the TensorFlow Model Optimization Toolkit (separate repo, tensorflow/model-optimization), but the kernels and grad ops live here in tensorflow/python/ops/.

After QAT, the converter can produce a more accurate INT8 TFLite model than post-training-only quantization.

MLIR tf_quant

tensorflow/compiler/mlir/quantization/ is a generic quantization framework on top of MLIR. It targets:

• TFLite tfl ops (the legacy mobile path).
• StableHLO operations (for portable deployments).
• Server-side TF graphs.

It supports per-axis weights, per-tensor activations, INT4/INT8/UINT8 schemes, weight-only quantization, and dynamic-range quantization.

Mixed precision

tf.keras.mixed_precision.set_global_policy("mixed_float16") enables Keras-level mixed precision: variables stay in float32, computations cast to float16, loss scaling is applied. Implementation in tensorflow/python/keras/mixed_precision/.

For non-Keras code, Grappler's auto_mixed_precision pass (tensorflow/core/grappler/optimizers/auto_mixed_precision.cc) can rewrite a graph to use float16 / bfloat16 where safe.

INT8 inference via TensorRT

tensorflow/compiler/tf2tensorrt/ integrates TensorRT, which has its own INT8 calibration. See compilers/tensorrt. The flow is similar: provide a representative dataset, calibrate, get a quantized engine wrapped in TRTEngineOp.

Integration points

• Converter — entry into TFLite quantization is tf.lite.TFLiteConverter.
• Grappler — auto_mixed_precision runs as a graph-level pass.
• Keras — mixed precision API.
• TF Model Optimization Toolkit — the recommended Python entry for QAT (separate repo).
• MLIR — virtually all new quantization work happens at the MLIR level.

Entry points for modification

• New TFLite quantization scheme — tensorflow/compiler/mlir/lite/quantization/ and tensorflow/lite/tools/optimize/.
• New mixed-precision rewrite — tensorflow/core/grappler/optimizers/auto_mixed_precision.cc.
• TensorRT INT8 changes — tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc.

Related

• apps/tensorflow-lite — the converter that consumes most quantization work.
• compilers/mlir — tf_quant dialect.
• compilers/tensorrt — INT8 via TRT.