This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fc21cd6ede [Fix][Relax][TFLite] Use astype for frontend casts (#19932)
fc21cd6ede is described below
commit fc21cd6ede10fe23b96644748b07b5e2a89cff43
Author: Hongyi Wu <[email protected]>
AuthorDate: Mon Jul 13 05:24:16 2026 +0800
[Fix][Relax][TFLite] Use astype for frontend casts (#19932)
## Summary
Fix TFLite Relax frontend cast paths that still used removed/nonexistent
cast
APIs.
- Use `relax.op.astype` for FLOAT16 `DEQUANTIZE` constants.
- Use `relax.op.astype` around the existing quantized `AVERAGE_POOL_2D`
converter path.
## Design
This PR only replaces invalid frontend API calls with `relax.op.astype`.
The quantized avgpool regression test calls the converter path directly
because
the top-level TFLite importer still rejects quantized `AVERAGE_POOL_2D`
before
conversion. Enabling that operator globally is out of scope.
## Tests
Added:
- `test_dequantize_float16_uses_astype`
- `test_quantized_avg_pool2d_uses_astype`
Validated with:
```bash
python -m ruff format \
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 \
-k "dequantize or avg_pool" -q
```
Result:
```text
ruff format: 2 files left unchanged
ruff check: All checks passed
targeted dequantize/avg_pool tests: 9 passed, 551 deselected
```
I also ran the full TFLite frontend file:
```text
tests/python/relax/test_frontend_tflite.py: 559 passed, 1 failed
```
The remaining failure is unrelated to this PR:
`test_broadcast_to` expects `R.multiply(..., ones)` while the importer
emits
`R.broadcast_to(...)`.
---
.../tvm/relax/frontend/tflite/tflite_frontend.py | 6 +-
tests/python/relax/test_frontend_tflite.py | 148 +++++++++++++++++++++
2 files changed, 151 insertions(+), 3 deletions(-)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 6918c600f6..a2db135466 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -5627,9 +5627,9 @@ class OperatorConverter:
"TFLite avg_pool2dreshape requires input and output scale"
"and zero points to be equal"
)
- out = relax.op.cast(in_expr, dtype="int32")
+ out = relax.op.astype(in_expr, "int32")
out = relax.op.nn.avg_pool2d(out, **params)
- out = relax.op.cast(out, dtype=output_tensor_type_str)
+ out = relax.op.astype(out, output_tensor_type_str)
else:
out = relax.op.nn.avg_pool2d(in_expr, **params)
elif pool_type == "max":
@@ -7327,7 +7327,7 @@ class OperatorConverter:
in_expr = self.exp_tab.new_const(
input_value, dtype=dtype,
source_name=input_tensor.tensor.Name()
)
- out = relax.cast(in_expr, dtype="float32")
+ out = relax.op.astype(in_expr, "float32")
return out
in_expr = self.get_expr(input_tensor.tensor_idx)
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index c756227fe0..dab0b4b103 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -4907,6 +4907,7 @@ _tfl_int32_vector =
_get_tflite_schema_module("Int32Vector")
_tfl_model = _get_tflite_schema_module("Model")
_tfl_operator = _get_tflite_schema_module("Operator")
_tfl_operator_code = _get_tflite_schema_module("OperatorCode")
+_tfl_pool2d_options = _get_tflite_schema_module("Pool2DOptions")
_tfl_quantization_parameters =
_get_tflite_schema_module("QuantizationParameters")
_tfl_sparsity_parameters = _get_tflite_schema_module("SparsityParameters")
_tfl_subgraph = _get_tflite_schema_module("SubGraph")
@@ -11188,6 +11189,153 @@ def test_dequantize_op_uses_relax_dequantize():
tvm.ir.assert_structural_equal(mod, Expected)
+def test_dequantize_float16_uses_astype():
+ """TFLite DEQUANTIZE float16 -> float32 uses R.astype."""
+ builder = flatbuffers.Builder(1024)
+
+ input_data = np.array([1.5, -2.0], dtype=np.float16)
+
+ input_tensor = _build_tensor(builder, 0, [2],
tensor_type=_tfl_tensor_type.FLOAT16)
+ output_tensor = _build_tensor(builder, 1, [2],
tensor_type=_tfl_tensor_type.FLOAT32)
+
+ dequantize_op = _build_operator(builder, 0, [0], [1])
+ subgraph = _build_subgraph(
+ builder,
+ tensors=[input_tensor, output_tensor],
+ operators=[dequantize_op],
+ inputs=[],
+ outputs=[1],
+ )
+ operator_codes = [_build_operator_code(builder,
_tfl_builtin_operator.DEQUANTIZE)]
+ input_buffer = _build_buffer(builder, input_data.tobytes())
+ output_buffer = _build_buffer(builder)
+ buf = _finish_tflite_model(
+ builder,
+ subgraph=subgraph,
+ operator_codes=operator_codes,
+ buffers=[input_buffer, output_buffer],
+ )
+
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+ mod = from_tflite(tflite_model)
+ mod["main"] = mod["main"].without_attr("params")
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((2,), dtype="float32"):
+ R.func_attr({"num_input": 0})
+ with R.dataflow():
+ gv: R.Tensor((2,), dtype="float32") = R.astype(
+ R.const(np.array([1.5, -2.0], dtype=np.float16)),
dtype="float32"
+ )
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_quantized_avg_pool2d_uses_astype():
+ """Quantized AVERAGE_POOL_2D casts through int32 with R.astype."""
+ builder = flatbuffers.Builder(1024)
+
+ qparams = _build_quantization_parameters(
+ builder, scale=[0.5], zero_point=[3], quantized_dimension=0
+ )
+ input_tensor = _build_tensor(
+ builder,
+ 0,
+ [1, 2, 2, 1],
+ tensor_type=_tfl_tensor_type.INT8,
+ quantization=qparams,
+ )
+ output_tensor = _build_tensor(
+ builder,
+ 1,
+ [1, 1, 1, 1],
+ tensor_type=_tfl_tensor_type.INT8,
+ quantization=qparams,
+ )
+
+ _tfl_pool2d_options.Pool2DOptionsStart(builder)
+ _tfl_pool2d_options.Pool2DOptionsAddPadding(builder, _tfl_padding.VALID)
+ _tfl_pool2d_options.Pool2DOptionsAddStrideH(builder, 1)
+ _tfl_pool2d_options.Pool2DOptionsAddStrideW(builder, 1)
+ _tfl_pool2d_options.Pool2DOptionsAddFilterHeight(builder, 2)
+ _tfl_pool2d_options.Pool2DOptionsAddFilterWidth(builder, 2)
+ _tfl_pool2d_options.Pool2DOptionsAddFusedActivationFunction(builder,
_tfl_activation_fn.NONE)
+ pool_opts = _tfl_pool2d_options.Pool2DOptionsEnd(builder)
+
+ avg_pool_op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1],
+ builtin_options_type=_tfl_builtin_options.Pool2DOptions,
+ builtin_options=pool_opts,
+ )
+ subgraph = _build_subgraph(
+ builder,
+ tensors=[input_tensor, output_tensor],
+ operators=[avg_pool_op],
+ inputs=[0],
+ outputs=[1],
+ )
+ operator_codes = [_build_operator_code(builder,
_tfl_builtin_operator.AVERAGE_POOL_2D)]
+ buf = _finish_tflite_model(
+ builder,
+ subgraph=subgraph,
+ operator_codes=operator_codes,
+ buffers=[_build_buffer(builder), _build_buffer(builder)],
+ )
+
+ if hasattr(tflite.Model, "Model"):
+ tflite_model = tflite.Model.Model.GetRootAsModel(buf, 0)
+ else:
+ tflite_model = tflite.Model.GetRootAsModel(buf, 0)
+
+ subgraph = tflite_model.Subgraphs(0)
+ bb = relax.BlockBuilder()
+ exp_tab = tflite_frontend.ExprTable()
+ input_var = relax.Var("tvmgen_tensor_0", relax.TensorType((1, 2, 2, 1),
dtype="int8"))
+ exp_tab.set_expr("tvmgen_tensor_0", input_var)
+ converter = tflite_frontend.OperatorConverter(tflite_model, subgraph,
exp_tab, bb)
+ with bb.function("main", [input_var]):
+ with bb.dataflow():
+ output = converter.convert_pool2d(subgraph.Operators(0), "average")
+ gv = bb.emit_output(output)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(tvmgen_tensor_0: R.Tensor((1, 2, 2, 1), dtype="int8")) ->
R.Tensor(
+ (1, 1, 1, 1), dtype="int8"
+ ):
+ with R.dataflow():
+ lv: R.Tensor((1, 2, 2, 1), dtype="int32") =
R.astype(tvmgen_tensor_0, dtype="int32")
+ lv1: R.Tensor((1, 1, 1, 1), dtype="int32") = R.nn.avg_pool2d(
+ lv,
+ pool_size=[2, 2],
+ strides=[1, 1],
+ dilation=[1, 1],
+ padding=[0, 0, 0, 0],
+ ceil_mode=False,
+ count_include_pad=False,
+ layout="NHWC",
+ out_layout="NHWC",
+ )
+ gv: R.Tensor((1, 1, 1, 1), dtype="int8") = R.astype(lv1,
dtype="int8")
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def test_quantized_conv2d_per_tensor_uses_qdq():
"""Quantized Conv2D with per-tensor quantization uses DQ -> conv2d -> Q."""
builder = flatbuffers.Builder(2048)
