Skip to content

[Relax][Frontend][TFLite] Add REDUCE_WINDOW support#19637

Merged
tlopex merged 1 commit into
apache:mainfrom
THINKER-ONLY:relax-tflite-reduce-window-v2
May 29, 2026
Merged

[Relax][Frontend][TFLite] Add REDUCE_WINDOW support#19637
tlopex merged 1 commit into
apache:mainfrom
THINKER-ONLY:relax-tflite-reduce-window-v2

Conversation

@THINKER-ONLY
Copy link
Copy Markdown
Contributor

Summary

Add Relax TFLite frontend support for the builtin REDUCE_WINDOW operator.
This covers the ordinary TFLite op only, not STABLEHLO_REDUCE_WINDOW.

The converter parses ReduceWindowOptions from BuiltinOptions2, validates
the static window attributes, and lowers supported reduce functions through
topi.sliding_window plus Relax reductions.

Supported modes:

  • ADD
  • MUL
  • MINIMUM
  • MAXIMUM
  • ALL
  • ANY

Empty output shapes are handled directly with relax.op.zeros. Quantized
REDUCE_WINDOW, dynamic window attributes, and unsupported reduce functions
remain rejected with explicit errors.

Testing

  • python -m py_compile 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 reduce_window -q -p no:tvm.testing.plugin
  • python -m pytest tests/python/relax/test_frontend_tflite.py -k "reduce_window or reduction_ops" -q -p no:tvm.testing.plugin
  • conda run -n test python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py

Related

Related to #19519.

Copilot AI review requested due to automatic review settings May 28, 2026 14:02
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for the REDUCE_WINDOW operator in the TFLite Relax frontend, along with comprehensive unit tests covering various reduction modes and error cases. The reviewer identified a critical issue where the call to topi.sliding_window in both the frontend implementation and the test helpers passes an incorrect number of arguments and incorrect argument types, which would lead to runtime failures. Correcting this call to use the proper 5-argument signature with explicit shape expressions for padding is recommended.

Comment on lines +3578 to +3587
windowed = relax.op.call_dps_packed(
"topi.sliding_window",
(
data,
0,
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(output_shape + dilated_window_shape, input_dtype),
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The topi.sliding_window packed function expects 5 arguments: data, pad_before, pad_after, window_shape, and strides. The current call only passes 4 arguments, omitting pad_after. Additionally, pad_before is passed as 0 (an integer) instead of a shape expression of zeros of length rank (e.g., relax.ShapeExpr([0] * rank)). This will cause a runtime failure when executing the compiled model.

Suggested change
windowed = relax.op.call_dps_packed(
"topi.sliding_window",
(
data,
0,
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(output_shape + dilated_window_shape, input_dtype),
)
windowed = relax.op.call_dps_packed(
"topi.sliding_window",
(
data,
relax.ShapeExpr([0] * rank),
relax.ShapeExpr([0] * rank),
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(output_shape + dilated_window_shape, input_dtype),
)

Comment on lines +4097 to +4110
windowed = bb.emit(
relax.op.call_dps_packed(
"topi.sliding_window",
(
x,
0,
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(
output_shape + tuple(dilated_window_shape), dtype
),
)
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the expected topi.sliding_window call in the test helper to match the corrected 5-argument signature with explicit pad_before and pad_after shape expressions.

            windowed = bb.emit(
                relax.op.call_dps_packed(
                    "topi.sliding_window",
                    (
                        x,
                        relax.ShapeExpr([0] * rank),
                        relax.ShapeExpr([0] * rank),
                        relax.ShapeExpr(dilated_window_shape),
                        relax.ShapeExpr(window_strides),
                    ),
                    out_sinfo=relax.TensorStructInfo(
                        output_shape + tuple(dilated_window_shape), dtype
                    ),
                )
            )

Comment on lines +4150 to +4163
windowed = bb.emit(
relax.op.call_dps_packed(
"topi.sliding_window",
(
x,
0,
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(
output_shape + tuple(dilated_window_shape), "bool"
),
)
)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the expected topi.sliding_window call in the boolean test helper to match the corrected 5-argument signature with explicit pad_before and pad_after shape expressions.

Suggested change
windowed = bb.emit(
relax.op.call_dps_packed(
"topi.sliding_window",
(
x,
0,
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(
output_shape + tuple(dilated_window_shape), "bool"
),
)
)
windowed = bb.emit(
relax.op.call_dps_packed(
"topi.sliding_window",
(
x,
relax.ShapeExpr([0] * rank),
relax.ShapeExpr([0] * rank),
relax.ShapeExpr(dilated_window_shape),
relax.ShapeExpr(window_strides),
),
out_sinfo=relax.TensorStructInfo(
output_shape + tuple(dilated_window_shape), "bool"
),
)
)

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds support for the TFLite REDUCE_WINDOW operator in the Relax TFLite frontend and introduces corresponding unit tests.

Changes:

  • Register and implement REDUCE_WINDOW conversion in the Relax TFLite frontend.
  • Add FlatBuffers test-model builders for REDUCE_WINDOW and expected Relax IR constructors.
  • Add coverage for supported modes and key failure cases (unsupported function, rank mismatch, invalid stride, shape mismatch, empty output).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File Description
tests/python/relax/test_frontend_tflite.py Adds schema bindings, model builder helpers, and tests for REDUCE_WINDOW.
python/tvm/relax/frontend/tflite/tflite_frontend.py Implements REDUCE_WINDOW conversion and wires it into the operator dispatch table.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

_tfl_operator = _get_tflite_schema_module("Operator")
_tfl_operator_code = _get_tflite_schema_module("OperatorCode")
_tfl_quantization_parameters = _get_tflite_schema_module("QuantizationParameters")
_tfl_reduce_window_options = _get_tflite_schema_module("ReduceWindowOptions")
_tfl_dimension_type = _get_tflite_schema_enum("DimensionType")
_tfl_fc_weights_format = _get_tflite_schema_enum("FullyConnectedOptionsWeightsFormat")
_tfl_padding = _get_tflite_schema_enum("Padding")
_tfl_reduce_window_function = _get_tflite_schema_enum("ReduceWindowFunction")
Comment on lines +3607 to +3608
raise tvm.error.OpNotImplemented(
f"TFLite REDUCE_WINDOW reduce_function {reduce_function} is not supported."
Comment on lines +3520 to +3522
window_shape = to_int_list(self.get_tensor_value(window_shape_tensor))
window_strides = to_int_list(self.get_tensor_value(window_strides_tensor))
window_dilations = to_int_list(self.get_tensor_value(window_dilations_tensor))
Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for adding this converter. I think this currently rejects valid TFLite REDUCE_WINDOW models because the init value is required to be rank-0 here.

The TFLite kernel tests build REDUCE_WINDOW with the init tensor as a one-element tensor with shape [1], not shape []. For example, the official test helper uses AddConstInput(kTensorType, {init_value_data_}, {1}). The op semantics only need a single init value, so accepting only len(shape) == 0 here is too strict and will reject models that the TFLite runtime accepts.

I think this should check that the init tensor has exactly one element instead of requiring rank 0, and then reshape/squeeze it to a scalar if needed before combining it with the reduced window result.

The current tests miss this because _build_reduce_window_model constructs the init tensor with shape []. Please add coverage for the [1] init tensor form, ideally with a numeric test matching TFLite behavior.

This commit adds Relax TFLite frontend support for the builtin REDUCE_WINDOW operator. The converter parses ReduceWindowOptions from BuiltinOptions2, validates the static window attributes, and lowers numeric and boolean reductions through topi.sliding_window plus Relax reductions.

The implementation covers ADD, MUL, MINIMUM, MAXIMUM, ALL, and ANY reduce functions. Empty output shapes are handled directly with Relax zeros, while quantized REDUCE_WINDOW and dynamic window attributes are left unsupported with explicit errors.

Tests add minimal hand-built TFLite flatbuffer fixtures and structural-equal coverage for all supported reduce functions, empty output dimensions, unsupported reduce functions, rank mismatch, and invalid stride values.
@THINKER-ONLY THINKER-ONLY force-pushed the relax-tflite-reduce-window-v2 branch from 3b3c8ce to 32e64fc Compare May 29, 2026 02:19
Copy link
Copy Markdown
Member

@tlopex tlopex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tlopex tlopex merged commit 0d70112 into apache:main May 29, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants