Aharrypotter opened a new pull request, #19491:
URL: https://github.com/apache/tvm/pull/19491
## Summary
This PR adds Relax TFLite frontend support for the following segment
operators from #19412:
- `SEGMENT_SUM`
- `UNSORTED_SEGMENT_MIN`
- `UNSORTED_SEGMENT_PROD`
These operators are lowered through `relax.op.scatter_nd` with the
corresponding reduction modes.
## Changes
### TFLite Frontend
1. Add TFLite converter mappings for segment operators:
- `SEGMENT_SUM` -> `scatter_nd(..., reduction="add")`
- `UNSORTED_SEGMENT_MIN` -> `scatter_nd(..., reduction="min")`
- `UNSORTED_SEGMENT_PROD` -> `scatter_nd(..., reduction="mul")`
2. Add shared segment lowering logic:
- Convert `segment_ids` into scatter indices via `expand_dims`.
- Build the output shape from `num_segments` or constant `segment_ids`.
- Initialize the scatter base tensor with the correct reduction
identity.
### Tests
Add TFLite frontend tests for:
- `test_segment_sum`
- `test_unsorted_segment_min`
- `test_unsorted_segment_prod`
Each test verifies the imported Relax IR lowers to `R.scatter_nd` with the
expected reduction mode and base tensor initialization.
## Testing
All targeted tests pass:
```bash
python -m pytest \
tests/python/relax/test_frontend_tflite.py::test_scatter_nd \
tests/python/relax/test_frontend_tflite.py::test_segment_sum \
tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_min \
tests/python/relax/test_frontend_tflite.py::test_unsorted_segment_prod \
-q
```
## References
- Issue #19412: TFLite Relax frontend operator support tracking
- Related PR #19490: Adds SCATTER_ND support
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
