[Relax][TFLite] Add gather frontend expected IRModule tests#19516
[Relax][TFLite] Add gather frontend expected IRModule tests#19516weicheng-hsu wants to merge 1 commit intoapache:mainfrom
Conversation
This adds explicit Expected IRModule coverage for TFLite GATHER and GATHER_ND frontend conversion. GATHER_ND uses Relax gather_nd with int64 indices, so the frontend now casts int32 TFLite indices to int64 before emitting the Relax op. This keeps the generated module well-typed and matches the expected Relax IR.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request modifies the TFLite frontend in Relax to cast gather_nd indices to int64 if they are provided as int32, satisfying the requirements of the Relax gather_nd operator. Additionally, new test cases for gather and gather_nd have been added. The reviewer suggests extending these tests to cover both int32 and int64 indices for both operators to ensure comprehensive coverage.
| verify(ReverseV2, Expected) | ||
|
|
||
|
|
||
| def test_gather(): |
There was a problem hiding this comment.
This test covers gather with int64 indices. For completeness, it would be beneficial to also add a test case for int32 indices, as TFLite's GATHER op supports both int32 and int64 for indices.
| verify(Gather, Expected) | ||
|
|
||
|
|
||
| def test_gather_nd(): |
There was a problem hiding this comment.
This test covers gather_nd with int32 indices, which is great for verifying the change in this PR. To make the test coverage for this operator more complete, could you also add a test case for int64 indices? This would ensure the int64 path is not broken by future changes and remains tested.
1 participant
This adds explicit Expected IRModule coverage for TFLite GATHER and GATHER_ND frontend conversion.
GATHER_ND uses Relax gather_nd with int64 indices, so the frontend now casts int32 TFLite indices to int64 before emitting the Relax op. This keeps the generated module well-typed and matches the expected Relax IR.
Testing:
python -m pytest tests/python/relax/test_frontend_tflite.py -k "gather"related to #18971