[Relax][Frontend][TFLite] Support StableHLO shape ops - #19869
Closed
Mohxen wants to merge 2 commits into
Closed
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 adds support for converting the STABLEHLO_RESHAPE operator to Relax in the TFLite frontend, including corresponding unit tests. The feedback suggests replacing the assert statements in the new conversion function with explicit ValueError checks to ensure validation is not bypassed when Python is run with optimization flags.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+3020
to
+3023
| input_tensors = self.get_input_tensors(op) | ||
| assert len(input_tensors) == 1 | ||
| output_tensors = self.get_output_tensors(op) | ||
| assert len(output_tensors) == 1 |
Contributor
There was a problem hiding this comment.
Using assert statements for input validation is discouraged because they can be optimized away when Python is run with the -O (optimize) flag. This would bypass the validation checks entirely and could lead to cryptic errors (like IndexError) later in the execution. It is safer and more robust to explicitly check the conditions and raise a ValueError with a descriptive error message.
Suggested change
| input_tensors = self.get_input_tensors(op) | |
| assert len(input_tensors) == 1 | |
| output_tensors = self.get_output_tensors(op) | |
| assert len(output_tensors) == 1 | |
| input_tensors = self.get_input_tensors(op) | |
| if len(input_tensors) != 1: | |
| raise ValueError(f"STABLEHLO_RESHAPE expects exactly 1 input tensor, but got {len(input_tensors)}") | |
| output_tensors = self.get_output_tensors(op) | |
| if len(output_tensors) != 1: | |
| raise ValueError(f"STABLEHLO_RESHAPE expects exactly 1 output tensor, but got {len(output_tensors)}") |
Mohxen
changed the title
74 tasks
tlopex
pushed a commit
that referenced
this pull request
Aug 18, 2026
## Summary This PR adds Relax TFLite frontend support for the remaining StableHLO shape operators tracked by #19519: - `STABLEHLO_RESHAPE` -> `R.reshape` - `STABLEHLO_SLICE` -> `R.strided_slice` - `STABLEHLO_TRANSPOSE` -> `R.permute_dims` It carries forward the implementation from #19869 by @Mohxen onto the current `main` branch and addresses the outstanding review feedback by using explicit `ValueError` checks for the input and output arity of all three new converters. ## Design ### StableHLO reshape `STABLEHLO_RESHAPE` has one tensor input and a statically described result shape. The converter reads that shape from the TFLite output tensor metadata and emits `relax.op.reshape`. ### StableHLO slice `STABLEHLO_SLICE` stores `start_indices`, `limit_indices`, and `strides` in `StablehloSliceOptions`. The converter parses those vectors, applies them to all input axes, and emits `relax.op.strided_slice`. ### StableHLO transpose `STABLEHLO_TRANSPOSE` stores its permutation in `StablehloTransposeOptions`. The converter parses the permutation and emits `relax.op.permute_dims`. ## Operator Support | Operator | TFLite metadata | Relax lowering | Supported subset | |---|---|---|---| | `STABLEHLO_RESHAPE` | output tensor shape | `R.reshape` | static result shape | | `STABLEHLO_SLICE` | start, limit, and stride vectors | `R.strided_slice` | static slice attributes | | `STABLEHLO_TRANSPOSE` | permutation vector | `R.permute_dims` | static permutation | ## Tests The tests manually build minimal TFLite flatbuffers for each StableHLO operator and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. The slice fixture exercises non-unit strides, and the transpose fixture uses a nontrivial three-dimensional permutation. 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 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 "stablehlo_reshape or stablehlo_slice or stablehlo_transpose" -q ``` Result: ```text ruff format --check: 2 files already formatted ruff check: All checks passed py_compile: passed targeted StableHLO shape tests: 3 passed, 551 deselected ``` ## References - Completes the remaining StableHLO shape-operator items in #19519. - Continues and supersedes #19869 by @Mohxen. --------- Co-authored-by: Mohxen <mohsenrahmati@icloud.com>
Member
|
Closed as supported by #20114 |
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
2 participants
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
Adds Relax TFLite frontend support for the remaining StableHLO shape operators:
STABLEHLO_RESHAPE->R.reshapeSTABLEHLO_SLICE->R.strided_sliceSTABLEHLO_TRANSPOSE->R.permute_dimsRelated to #19519.
Testing
python -m ruff check python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.pypython -m py_compile python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.pygit diff --checkRelated to #19519.