[Relax][TFLite] Fix and test DEPTH_TO_SPACE/SPACE_TO_DEPTH, SELECT ops#19381
Conversation
…tests for SELECT, DEPTH_TO_SPACE, SPACE_TO_DEPTH
| block_size = space_to_depth_options.BlockSize() | ||
| out = relax.op.nn.space_to_depth(in_expr, block_size, layout="NHWC") | ||
|
|
||
| # TFLite uses NHWC layout: (N, H, W, C) -> (N, H/bs, W/bs, C*bs*bs) |
There was a problem hiding this comment.
convert_depth_to_space and convert_space_to_depth in tflite_frontend.py call
relax.op.nn.depth_to_space / space_to_depth, but these ops don't exist in Relax,
causing AttributeError at runtime.
![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 updates the TFLite frontend for TVM Relax, specifically reimplementing the depth_to_space and space_to_depth operators. Instead of using the high-level Relax operations, these are now implemented using a sequence of reshape and permute_dims to explicitly handle the NHWC layout used by TFLite. Additionally, the PR includes several formatting improvements across the frontend and test suite, and introduces new unit tests for the select, depth_to_space, and space_to_depth operators. I have no feedback to provide as there were no review comments to evaluate.
tlopex
left a comment
tlopex
left a comment
There was a problem hiding this comment.
LGTM. Thank you for the contribution!
2 participants
Summary
part of: #18971
DEPTH_TO_SPACEandSPACE_TO_DEPTHconverters in the TFLite Relax frontend that were calling non-existentrelax.op.nn.depth_to_space/relax.op.nn.space_to_depth, causingAttributeErrorat runtimeSELECT/WHERE,DEPTH_TO_SPACE, andSPACE_TO_DEPTHoperators