Description
This PR adds support for the CONV_3D operator in the TFLite frontend for
Relax.
Key Changes
- Operator Mapping: Added CONV_3D to the OperatorConverter mapping in
tflite_frontend.py.
- Implementation:
- Implemented convert_conv3d to handle 3D convolution attributes such as
StrideD/H/W, DilationD/H/W, and Padding.
- Correctly handled the TFLite 3D kernel layout, which is expected to be
DHWIO (Depth, Height, Width, Input Channels, Output Channels).
- Integrated support for fused activation functions (ReLU, ReLU6, etc.)
directly following the convolution.
- Unit Tests:
- Added comprehensive tests in
tests/python/relax/test_frontend_tflite.py covering:
- VALID and SAME padding modes.
- Various stride and dilation configurations.
- Verification against expected Relax IR structure.
Testing:
- `python3 -m pytest tests/python/relax/test_frontend_tflite.py -k
"test_conv3d"`
Notes for Reviewers
The implementation follows the existing pattern used for CONV_2D but
extends it to the 5D case (NDHWC layout). I've ensured that the kernel
layout mapping aligns with TVM's R.nn.conv3d requirements.
Related to: https://github.com/apache/tvm/issues/19519
Wei-Cheng Hsu
4ab312bf12