[Relax][Frontend][TFLite] Add support for FFT/complex operators: REAL, IMAG, COMPLEX_ABS, RFFT2D#19763
[Relax][Frontend][TFLite] Add support for FFT/complex operators: REAL, IMAG, COMPLEX_ABS, RFFT2D#19763fnhirwa wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request adds support for converting TFLite REAL, IMAG, and COMPLEX_ABS operators to TVM Relax, representing complex64 tensors as float32 tensors with an extra trailing dimension of size 2. It also adds a placeholder for RFFT2D which raises an unimplemented error. The review feedback suggests using negative indexing (axes=[-1] and axis=[-1]) instead of calculating the last axis dynamically using the tensor's ndim to prevent potential errors with dynamic or unknown ranks, along with corresponding updates to the test assertions.
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.
tlopex
left a comment
There was a problem hiding this comment.
Thanks for working on this. I think this needs a small fix before merge.
The complex64 -> float32[..., 2] lowering is currently applied in two different places, which can produce wrong input shapes when callers pass shape_dict or dtype_dict overrides. _input_type() already converts complex inputs to float32[..., 2], then from_tflite() applies user overrides and may append another trailing 2 if the overridden dtype is complex64.
For example:
dtype_dict={"x": "complex64"}can turn an already packed(2, 4, 2)shape into(2, 4, 2, 2).shape_dict={"x": (2, 4)}overwrites the packed default shape while the dtype remainsfloat32, soREAL/IMAGslice the wrong last axis.
I think the safer structure is to keep _input_type() returning the logical TFLite shape/dtype, apply user overrides, and then lower complex64 to float32[..., 2] exactly once.
A couple of related scope/test issues:
RFFT2Dis registered inconvert_map, butconvert_rfft2d()always raisesOpNotImplemented. The PR title/body says RFFT2D support and mentions an RFFT2D test, but the patch only adds tests forREAL,IMAG, andCOMPLEX_ABS. Please either implement it, or keep it clearly unsupported and add a negative test / update the PR description.- The new complex representation is only wired for model inputs. Constant or non-input complex tensors still go through the generic tensor path, which does not decode
TensorType.COMPLEX64. If input-only support is intentional for this PR, please make that scope explicit and cover it with tests.
Part of #19519
This PR adds support for the FFT and complex operator family in the Relax TFLite frontend.
Key implementations:
REAL,IMAG,COMPLEX_ABS, andRFFT2Dto the TFLite op map.convert_realandconvert_imagwhich extract the real and imaginary parts of a complex tensor viastrided_slice+squeezealong the last axis.convert_complex_abswhich computessqrt(re^2 + im^2)using elementwise Relax ops.complex64tensors (which have no native Relax dtype equivalent) are represented asfloat32[..., 2], where the last axis holds(real, imaginary)interleaved.What is missing:
RFFT2DraisesOpNotImplementedin this PR. A nativerelax.op.signal.rfft2dop with a C++ registered backend is required.topi.signal.dftexists as pure Python TE but has noTVM_REGISTER_GLOBALentry and cannot be called viacall_dps_packed.Testing:
REAL,IMAG, andCOMPLEX_ABSintest_frontend_tflite.pyfollowing theverify(TestClass, Expected)pattern.RFFT2Dtest covers frontend conversion correctness.python3 -m pytest tests/python/relax/test_frontend_tflite.py -k "test_real or test_imag or test_complex_abs"