[Relax][Frontend][TFLite] Add REDUCE_WINDOW support#19556
Open
THINKER-ONLY wants to merge 1 commit into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements the REDUCE_WINDOW operator in the TFLite frontend for TVM Relax, supporting both numeric (ADD, MUL, MIN, MAX) and boolean (ALL, ANY) reduction functions. The implementation utilizes topi.sliding_window for the core windowing logic and includes comprehensive validation for input shapes, dtypes, and constant window parameters. Corresponding unit tests have been added to verify various modes and error conditions. Feedback suggests refactoring the series of conditional statements in the reduction logic into a dictionary mapping to improve code maintainability and readability.
Comment on lines
+2608
to
+2621
| if reduce_function == ReduceWindowFunction.ADD: | ||
| return relax.op.add(relax.op.sum(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MUL: | ||
| return relax.op.multiply(relax.op.prod(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MINIMUM: | ||
| return relax.op.minimum(relax.op.min(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MAXIMUM: | ||
| return relax.op.maximum(relax.op.max(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.ALL: | ||
| reduced = relax.op.min(relax.op.astype(windowed, "int8"), axis=reduce_axes) | ||
| return relax.op.logical_and(relax.op.astype(reduced, "bool"), init_value) | ||
| if reduce_function == ReduceWindowFunction.ANY: | ||
| reduced = relax.op.max(relax.op.astype(windowed, "int8"), axis=reduce_axes) | ||
| return relax.op.logical_or(relax.op.astype(reduced, "bool"), init_value) |
Contributor
There was a problem hiding this comment.
This series of if statements can be refactored using dictionaries to map reduction functions to their corresponding Relax operators. This would improve readability and maintainability by grouping the numeric and boolean reduction logic separately.
Suggested change
| if reduce_function == ReduceWindowFunction.ADD: | |
| return relax.op.add(relax.op.sum(windowed, axis=reduce_axes), init_value) | |
| if reduce_function == ReduceWindowFunction.MUL: | |
| return relax.op.multiply(relax.op.prod(windowed, axis=reduce_axes), init_value) | |
| if reduce_function == ReduceWindowFunction.MINIMUM: | |
| return relax.op.minimum(relax.op.min(windowed, axis=reduce_axes), init_value) | |
| if reduce_function == ReduceWindowFunction.MAXIMUM: | |
| return relax.op.maximum(relax.op.max(windowed, axis=reduce_axes), init_value) | |
| if reduce_function == ReduceWindowFunction.ALL: | |
| reduced = relax.op.min(relax.op.astype(windowed, "int8"), axis=reduce_axes) | |
| return relax.op.logical_and(relax.op.astype(reduced, "bool"), init_value) | |
| if reduce_function == ReduceWindowFunction.ANY: | |
| reduced = relax.op.max(relax.op.astype(windowed, "int8"), axis=reduce_axes) | |
| return relax.op.logical_or(relax.op.astype(reduced, "bool"), init_value) | |
| _NUMERIC_REDUCE_MAP = { | |
| ReduceWindowFunction.ADD: (relax.op.sum, relax.op.add), | |
| ReduceWindowFunction.MUL: (relax.op.prod, relax.op.multiply), | |
| ReduceWindowFunction.MINIMUM: (relax.op.min, relax.op.minimum), | |
| ReduceWindowFunction.MAXIMUM: (relax.op.max, relax.op.maximum), | |
| } | |
| if reduce_function in _NUMERIC_REDUCE_MAP: | |
| reduce_op, combine_op = _NUMERIC_REDUCE_MAP[reduce_function] | |
| return combine_op(reduce_op(windowed, axis=reduce_axes), init_value) | |
| _BOOL_REDUCE_MAP = { | |
| ReduceWindowFunction.ALL: (relax.op.min, relax.op.logical_and), | |
| ReduceWindowFunction.ANY: (relax.op.max, relax.op.logical_or), | |
| } | |
| if reduce_function in _BOOL_REDUCE_MAP: | |
| reduce_op, combine_op = _BOOL_REDUCE_MAP[reduce_function] | |
| reduced = reduce_op(relax.op.astype(windowed, "int8"), axis=reduce_axes) | |
| return combine_op(relax.op.astype(reduced, "bool"), init_value) |
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds Relax TFLite frontend support for the builtin REDUCE_WINDOW operator, including parsing ReduceWindowOptions from BuiltinOptions2 and lowering supported reduce functions via topi.sliding_window plus Relax reductions.
Changes:
- Adds
REDUCE_WINDOWconversion in the Relax TFLite frontend with validation for static window attributes and supported reduction modes. - Adds TFLite model builders + expected Relax IR constructors to test numeric, boolean, empty-output, and error cases for
REDUCE_WINDOW.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
python/tvm/relax/frontend/tflite/tflite_frontend.py |
Implements REDUCE_WINDOW lowering and validation in the TFLite-to-Relax converter. |
tests/python/relax/test_frontend_tflite.py |
Adds targeted unit tests and FlatBuffer model builders to cover supported/unsupported REDUCE_WINDOW behaviors. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+2475
to
+2476
| assert len(input_tensors) == 5, "input tensors length should be 5" | ||
| assert len(output_tensors) == 1, "output tensors length should be 1" |
Comment on lines
+2501
to
+2502
| input_shape = to_int_list(self.get_tensor_shape(input_tensor)) | ||
| output_shape = to_int_list(self.get_tensor_shape(output_tensor)) |
| (window_dim - 1) * dilation + 1 | ||
| for window_dim, dilation in zip(window_shape, window_dilations) | ||
| ] | ||
|
|
Comment on lines
+2576
to
+2581
| if any( | ||
| input_dim < dilated_dim | ||
| for input_dim, dilated_dim in zip(input_shape, dilated_window_shape) | ||
| ): | ||
| raise tvm.error.OpNotImplemented( | ||
| "TFLite REDUCE_WINDOW input/output shapes are inconsistent." |
Comment on lines
+2608
to
+2622
| if reduce_function == ReduceWindowFunction.ADD: | ||
| return relax.op.add(relax.op.sum(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MUL: | ||
| return relax.op.multiply(relax.op.prod(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MINIMUM: | ||
| return relax.op.minimum(relax.op.min(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.MAXIMUM: | ||
| return relax.op.maximum(relax.op.max(windowed, axis=reduce_axes), init_value) | ||
| if reduce_function == ReduceWindowFunction.ALL: | ||
| reduced = relax.op.min(relax.op.astype(windowed, "int8"), axis=reduce_axes) | ||
| return relax.op.logical_and(relax.op.astype(reduced, "bool"), init_value) | ||
| if reduce_function == ReduceWindowFunction.ANY: | ||
| reduced = relax.op.max(relax.op.astype(windowed, "int8"), axis=reduce_axes) | ||
| return relax.op.logical_or(relax.op.astype(reduced, "bool"), init_value) | ||
|
|
THINKER-ONLY
force-pushed
the
relax-tflite-reduce-window
branch
from
e75eebc to
468f2ef
Compare
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
force-pushed
the
relax-tflite-reduce-window
branch
from
468f2ef to
40d0cff
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add Relax TFLite frontend support for the builtin
REDUCE_WINDOWoperator.This covers the ordinary TFLite op only, not
STABLEHLO_REDUCE_WINDOW.The converter parses
ReduceWindowOptionsfromBuiltinOptions2, validatesthe static window attributes, and lowers supported reduce functions through
topi.sliding_windowplus Relax reductions.Supported modes:
ADDMULMINIMUMMAXIMUMALLANYEmpty output shapes are handled directly with
relax.op.zeros. QuantizedREDUCE_WINDOW, dynamic window attributes, and unsupported reduce functionsremain rejected with explicit errors.
Testing
python -m py_compile python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.pypython -m pytest tests/python/relax/test_frontend_tflite.py -k reduce_window -q -p no:tvm.testing.pluginpython -m pytest tests/python/relax/test_frontend_tflite.py -k "reduce_window or reduction_ops" -q -p no:tvm.testing.pluginconda run -n test python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.pyRelated
Related to #19519.