## Summary
This PR adds initial quantized TFLite import support to the Relax
frontend by
preserving tensor quantization metadata and replacing placeholder
`_qnn.op.*`
frontend calls with an explicit QDQ decomposition:
```text
dequantize -> float Relax op -> quantize
```
Before this PR, the Relax TFLite frontend raised `NotImplementedError`
as soon
as quantization metadata was seen during tensor parsing. This made
quantized
TFLite models unreachable. This PR keeps `scale`, `zero_point`, and
`QuantizedDimension()` in `TensorWrapper.qnn_params`, then uses the
existing
`R.quantize` / `R.dequantize` operators to lower supported quantized
paths.
The previous `_qnn.op.*` paths were effectively unreachable for normal
quantized TFLite models because `get_tensors()` raised
`NotImplementedError`
as soon as valid quantization metadata was parsed. After removing that
blocker,
those paths also needed to be replaced because they depended on
undefined
`_qnn` helpers and did not handle Relax QDQ, axis remapping, or
quantized bias
consistently.
Closes#19534.
## Design
Relax already has `R.quantize` and `R.dequantize` with C++ registration,
Python
APIs, legalization, and tests. Instead of introducing new fused Relax
QNN ops
for this first import PR, the frontend now decomposes quantized TFLite
operators
through QDQ around ordinary Relax float operators.
This keeps the change scoped to the Python TFLite frontend and existing
Relax
QDQ operators, while establishing a working import path first. Fused
int8 Relax
QNN operators can still be considered later if backend kernel selection
requires
them.
## Updated Converters
| Converter | Replacement |
|---|---|
| `get_tensors` | Preserve `scale`, `zero_point`, and
`QuantizedDimension()` |
| `quantize` / `dequantize` helpers | Use `R.quantize` / `R.dequantize`
with `axis` |
| `convert_quantize` | `float -> Q` and quantized requantize as `DQ ->
Q` |
| `convert_dequantize` | Use `R.dequantize` |
| `convert_relu`, `convert_relu6`, `convert_relu_n1_to_1` | `DQ ->
activation -> Q` |
| `_convert_elemwise` | Quantized binary ops use `DQ -> op -> fused
activation -> Q`; comparisons use `DQ -> compare` |
| `convert_reshape` | uint8 different-qparams path uses `DQ -> reshape
-> Q` |
| `_convert_reduce` | Quantized reduce uses `DQ -> reduce -> Q` |
| `convert_conv` | Quantized Conv2D uses `DQ input + DQ weight -> conv2d
-> Q` |
| `convert_fully_connected` | Quantized FC uses `DQ input + DQ weight ->
matmul -> Q` |
| `convert_concatenation` | Quantized concat uses `DQ each -> concat ->
Q` |
| `convert_transpose_conv` | Quantized transpose conv uses `DQ input +
DQ weight -> conv2d_transpose -> Q` |
| `convert_detection_postprocess` | Inline `_qnn.op.dequantize` calls
replaced with `self.dequantize` |
All `_qnn.op.*` references are removed, and the stale `# ruff: noqa:
F821`
suppression is no longer needed.
## Axis Remapping
The most correctness-sensitive part of this PR is axis remapping for
per-channel
weight dequantization after the frontend rewrites TFLite layouts into
Relax
layouts.
| Op | TFLite layout | Relax layout | Axis remap |
|---|---|---|---|
| Conv2D | `[OC, KH, KW, IC]` | `[KH, KW, IC, OC]` (`HWIO`) | `0 -> 3` |
| FullyConnected | `[OC, IC]` | `[IC, OC]` | `0 -> 1` |
| TransposeConv | `[OC, KH, KW, IC]` (`OHWI`) | `[IC, OC, KH, KW]`
(`IOHW`) | `0 -> 1` |
| DepthwiseConv | `[1, KH, KW, C*M]` | `[KH, KW, C, M]` (`HWOI`) |
per-channel unsupported |
For Conv2D, FC, and TransposeConv, non-zero weight
`QuantizedDimension()` values
are rejected with `OpAttributeInvalid`, because the supported quantized
TFLite
weight layout uses output-channel axis 0.
Per-channel depthwise convolution is guarded with `OpNotImplemented`.
The
TFLite depthwise reshape changes the channel-axis semantics in a way
that this
initial QDQ lowering does not represent directly.
## Bias Handling
TFLite INT32/INT64 bias tensors may not store explicit quantization
metadata.
For quantized Conv2D, FullyConnected, and TransposeConv, the frontend
follows
the implicit TFLite convention and dequantizes integer bias using:
```text
bias_scale = input_scale * weight_scale
bias_zero_point = 0
axis = 0
```
This supports both per-tensor and per-channel weight scales. The
per-channel
case is covered by a structural regression test that expects vector bias
scale.
## Fused Activation Handling
Conv2D, FullyConnected, and quantized concat preserve the existing
quantized-domain fused activation behavior:
```text
float op -> Q -> quantized-domain clip
```
The elemwise QDQ path applies fused activation before the final
quantize:
```text
DQ -> float binary op -> float fused activation -> Q
```
Both paths are intentional and covered by regression tests:
- quantized concat fused `RELU` checks the quantized-domain clip path
- quantized add fused `RELU6` checks the float-domain
activation-before-Q path
This PR also fixes a latent `R.clip` call-site bug in the quantized
fused
`RELU` helper by using `max=` rather than the unsupported `a_max=`
keyword.
## Safety Checks
- Quantized elemwise non-comparison outputs must have output qparams.
Missing
output quantization metadata now raises `OpAttributeInvalid` instead of
silently returning a float result.
- Per-channel quantization rejects non-zero per-axis zero points,
following the
TFLite quantization specification.
- Per-channel depthwise convolution is explicitly unsupported rather
than
importing with an incorrect axis interpretation.
## Tests
The new tests build minimal TFLite flatbuffers directly and compare the
imported
Relax IR with `tvm.ir.assert_structural_equal`. Unsupported-boundary
tests use
`pytest.raises`.
The FlatBuffer tests use schema module helpers instead of top-level
generated
builder functions when needed, so they work with the `tflite` Python
package
available in CI.
| Test | Coverage |
|---|---|
| `test_tensor_quantization_parameters_are_parsed` | per-tensor and
per-axis metadata parsing |
| `test_quantize_op_uses_relax_quantize` | TFLite `QUANTIZE` float input
|
| `test_quantize_op_requantize_uses_dq_q` | TFLite `QUANTIZE` as
requantize |
| `test_dequantize_op_uses_relax_dequantize` | TFLite `DEQUANTIZE` |
| `test_quantized_add_uses_qdq` | quantized ADD with differing input
qparams |
| `test_quantized_add_fused_relu6_uses_float_clip_before_quantize` |
elemwise fused activation before Q |
| `test_quantized_add_without_output_qparams_invalid` | invalid missing
output qparams guard |
| `test_quantized_conv2d_per_tensor_uses_qdq` | Conv2D per-tensor QDQ |
| `test_quantized_conv2d_per_channel_weight_uses_remapped_axis` | Conv2D
per-channel weight axis `0 -> 3` |
| `test_quantized_conv2d_with_int32_bias_dequantizes_bias` | Conv2D
INT32 bias scale |
|
`test_quantized_conv2d_per_channel_weight_with_int32_bias_dequantizes_bias`
| Conv2D per-channel vector bias scale |
| `test_quantized_concat_uses_qdq` | concat QDQ path |
| `test_quantized_concat_fused_relu_uses_quantized_clip` |
quantized-domain fused RELU clip |
| `test_per_channel_depthwise_conv_unsupported` | per-channel depthwise
guard |
| `test_uint8_reshape_requantize_uses_dq_reshape_q` | uint8 reshape with
different qparams |
| `test_transpose_conv_with_int32_bias_dequantizes_bias` | TransposeConv
INT32 bias DQ |
| `test_quantized_fully_connected_with_int32_bias_dequantizes_bias` | FC
INT32 bias DQ |
Local validation:
```bash
python -m ruff format --check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest tests/python/relax/test_frontend_tflite.py -q
```
Result:
```text
433 passed
```
## Limitations
- This PR prioritizes correct import and explicit Relax IR over fused
int8
kernel selection. The generated IR uses QDQ and float Relax operators.
- Per-channel depthwise convolution remains unsupported.
- The tests are structural IR tests. Numerical comparison against TFLite
runtime
outputs is left to follow-up work.
## References
- Issue #19534: Support quantized TFLite import in Relax frontend
- TFLite quantization spec:
https://www.tensorflow.org/lite/performance/quantization_spec
HoYi
2441461d12