[Fix][Relax][Frontend][TFLite] Tflite quantized avgpool by Theoo1997 · Pull Request #20298 · apache/tvm

Theoo1997 · 2026-09-09T19:49:28Z

Problem. AVERAGE_POOL_2D was missing from _SUPPORTED_QUANTIZED_OPS, so quantized models failed to import. Once reachable, the lowering was numerically wrong: TFLite's quantized AveragePool rounds the window average half away from zero, while relax's integer avg_pool2d truncates — a systematic bias, wrong on about half of all inputs.

Fix. Add the op to the allowlist. nn.avg_pool2d is unchanged, since altering a general operator's rounding would affect every integer user; instead, the frontend takes the window sum (pre-scaling makes the pool's own division exact) and applies TFLite's rounding explicitly.

Validation. Updated the existing structural test and added a numeric one; both fail before this change and pass after.

_SUPPORTED_QUANTIZED_OPS is checked before dispatch, and AVERAGE_POOL_2D was missing from it even though convert_pool2d already has a complete quantized branch for pool_type="average". The op was implemented and unreachable at the same time, so importing any full-integer-quantized classifier with a global average pool -- the standard head of MobileNet/ResNet/EfficientNet -- failed with OpNotImplemented. Verified by importing the MLCommons Tiny quantized CIFAR-10 ResNet, which now lowers to the expected integer pool: lv73 = R.astype(lv72, "int32") lv74 = R.nn.avg_pool2d(lv73, pool_size=[8, 8], strides=[8, 8], ...) lv75 = R.astype(lv74, "int8") Also adds TFLITE_QUANTIZED_SUPPORT.md, which documents this alongside a separate and still-open numerical bug found while validating the import: integer nn.avg_pool2d truncates the division where TFLite's quantized AveragePool rounds half away from zero, which is wrong on 537/1024 pooled values on this model and reaches the logits as 82/320 differing by up to 91. That one is described, measured and left unfixed -- the fix changes a general operator's semantics and is TVM's call to make.

TFLite's quantized AveragePool rounds the window average HALF AWAY FROM ZERO (reference/integer_ops/pooling.h): acc = acc > 0 ? (acc + count / 2) / count : (acc - count / 2) / count relax's integer nn.avg_pool2d divides with a truncating division, so every pooled value was biased toward zero by up to half an LSB. That is a systematic bias, not a rounding tie: on the MLCommons Tiny quantized ResNet it is wrong on 510 of 1024 pooled values, and because such graphs end in an int8 softmax it reaches the logits as 312/1280 differing by up to 106. nn.avg_pool2d itself is left alone -- changing the rounding of a general operator would change semantics for every integer user. Instead the frontend takes the window sum and performs TFLite's division explicitly. Pre-scaling the input by the window size makes avg_pool2d's own division exact, so what comes back is the sum; count_include_pad=True keeps that divisor constant (the padded taps are zeros and contribute nothing); relax.op.divide on int32 truncates toward zero, which is the semantics TFLite's formula is written against; and the per-position count of non-padded taps is folded to a constant at import, since the shapes are static. Measured with the new apps/tflite_quantized/verify_quantized_tflite.py, TVM default lowering, no BYOC: isolated AVERAGE_POOL_2D whole model before 510/1024 wrong (50%) 312/1280, max 106 after 0/1024 exact 22/1280, max 23 The whole-model figure is not zero for a separate reason, noted in the doc: TVM's QDQ lowering accumulates convolutions in float32. The script therefore also slices the operator out into a standalone one-op model and compares that on its own, which is the exact test, and keys its exit status on it.

…D rounding Moves the verification out of apps/ and into the suite where the rest of the TFLite frontend is tested, per convention: * test_quantized_avg_pool2d_uses_astype is renamed to ..._rounds_half_away_from_zero and updated for the new lowering, which it otherwise fails against. * A numeric test is added that pins the ANSWER rather than the shape, against the reference formula in tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h, on window sums that land on the .5 boundaries where truncation and round-half-away disagree. This is the test that would have caught the original bug: a truncating divide is wrong on about half of all inputs while looking perfectly reasonable in the IR. Both tests fail against the truncating lowering and pass with the fix. Also corrects the comment on the int32 pin after the pool. The widening is done by relax's own type inference for integer avg_pool2d when the window is large enough that an int32 accumulator could overflow -- an 8x8 pool over int32 comes back as int64, a 2x2 one stays int32 -- not by legalization as previously stated.

tlopex

The overflow check assumes 8-bit input (255 * window²), but this branch also accepts quantized int16. For a (1,17,17,1) input filled with 32767, a VALID 17×17 average pool returns -18657 instead of TFLite’s 32767: the pre-scaled sum overflows the int32 accumulator.

Could you use a sufficiently wide accumulator or sum the raw values directly？

…he input type The exact-sum lowering pre-scales the input by the window size, so its accumulator must hold max|x| * window^2. The overflow guard assumed 8-bit input (255 * window^2), but the branch also accepts int16: a VALID 17x17 pool of int16 32767 needs 2.7e9 and returned -18657 instead of 32767. Derive the bound from the input dtype's range (128 int8, 255 uint8, 32768 int16) and use int32 when it fits, int64 otherwise; the rounding and the divide run in the same type. int8 lowering is unchanged. Adds test_quantized_avg_pool2d_large_window_does_not_overflow (int16 +-extremes, int8, uint8, 17x17 window).

Theoo1997 · 2026-09-11T05:12:21Z

Thanks I confirme the bug. Rith int16 input a 1717 pool overflowed the int32 accumulator (32767 -> -18657). The accumulator width is now derived from the input dtype's range (max(|x|) * window^2): int32 when that fits, int64 otherwise, and the rounding/divide run in the same type. int8 lowering is unchanged. Added test_quantized_avg_pool2d_large_window_does_not_overflow covering int16 ±extremes, int8 and uint8 with a 1717 window.

tlopex · 2026-09-12T00:58:35Z

The divisor tensor is built from the serialized output shape, which becomes stale when from_tflite(..., shape_dict=...) overrides the input dimensions.

For an int8 model with input (1,4,4,1), a 2×2 VALID average pool, and stride 2:

Overriding the input to (1,2,2,1) should produce (1,1,1,1), but the old counts tensor silently broadcasts the result to (1,2,2,1).
Overriding it to (1,6,6,1) fails during import because the actual 3×3 pooled output cannot broadcast with the 2×2 counts tensor.

Could you derive the counts from the actual Relax input and pooled output shapes?

… Relax shapes The per-position divisor (and the SAME padding) were folded to constants from the SERIALIZED TFLite tensor shapes, which are stale once from_tflite(..., shape_dict=...) overrides the input dimensions. For an int8 model with input (1,4,4,1), a 2x2 VALID pool and stride 2, overriding the input to (1,2,2,1) silently broadcast the result back to (1,2,2,1) instead of (1,1,1,1), and (1,6,6,1) failed to import because the actual 3x3 pooled output could not broadcast with the 2x2 counts tensor. Read the input H/W from the Relax input expression and the output H/W from the pooled Relax tensor. The input extents also drive the SAME padding, so that is now correct for overridden shapes on every Pool2D path (average, max, L2). Non-static shapes raise OpAttributeUnImplemented. Adds test_quantized_avg_pool2d_follows_shape_dict_override: the two VALID overrides from the review plus two SAME overrides whose border windows see fewer taps, checked for shape and against a NumPy reference of TFLite's integer average pool.

Theoo1997 · 2026-09-15T14:20:46Z

The divisor tensor is built from the serialized output shape, which becomes stale when from_tflite(..., shape_dict=...) overrides the input dimensions.

For an int8 model with input (1,4,4,1), a 2×2 VALID average pool, and stride 2:

Overriding the input to (1,2,2,1) should produce (1,1,1,1), but the old counts tensor silently broadcasts the result to (1,2,2,1).

Overriding it to (1,6,6,1) fails during import because the actual 3×3 pooled output cannot broadcast with the 2×2 counts tensor.

Could you derive the counts from the actual Relax input and pooled output shapes?

I have derived the counts from the actual Relax input. Please check the new PR and let me know I need to do anythink else !!

tlopex · 2026-09-15T18:33:21Z

Could you convert only H and W to int here? Converting all four dimensions breaks imports with symbolic batch or channel dimensions that previously worked. Please also add a regression test with a symbolic batch dimension.

Theoo1997 and others added 6 commits September 9, 2026 18:09

Delete TFLITE_QUANTIZED_SUPPORT.md

f6ee904

Update test_frontend_tflite.py

67058e8

Update test_frontend_tflite.py

68cd869

Theoo1997 changed the title ~~Tflite quantized avgpool~~ [Fix][Relax][Frontend][TFLite] Tflite quantized avgpool Sep 10, 2026

tlopex requested changes Sep 10, 2026

View reviewed changes

Theoo1997 force-pushed the tflite-quantized-avgpool branch from c3614c2 to a2adbb2 Compare September 11, 2026 05:08

Theoo1997 closed this Sep 15, 2026

Theoo1997 reopened this Sep 15, 2026

Navigation Menu

Uh oh!

[Fix][Relax][Frontend][TFLite] Tflite quantized avgpool - #20298

[Fix][Relax][Frontend][TFLite] Tflite quantized avgpool#20298
Theoo1997 wants to merge 8 commits into
apache:mainfrom
Theoo1997:tflite-quantized-avgpool

Theoo1997 commented Sep 9, 2026

Uh oh!

tlopex left a comment

Uh oh!

Theoo1997 commented Sep 11, 2026

Uh oh!

tlopex commented Sep 12, 2026

Uh oh!

Theoo1997 commented Sep 15, 2026 •

edited

Loading

Uh oh!

tlopex commented Sep 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Theoo1997 commented Sep 9, 2026

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

Theoo1997 commented Sep 11, 2026

Uh oh!

tlopex commented Sep 12, 2026

Uh oh!

Theoo1997 commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tlopex commented Sep 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Theoo1997 commented Sep 15, 2026 •

edited

Loading