Skip to content

[Relax][Frontend][TFLite] Add support for FFT/complex operators: REAL, IMAG, COMPLEX_ABS, RFFT2D#19763

Open
fnhirwa wants to merge 2 commits into
apache:mainfrom
fnhirwa:fft
Open

[Relax][Frontend][TFLite] Add support for FFT/complex operators: REAL, IMAG, COMPLEX_ABS, RFFT2D#19763
fnhirwa wants to merge 2 commits into
apache:mainfrom
fnhirwa:fft

Conversation

@fnhirwa

@fnhirwa fnhirwa commented Jun 14, 2026

Copy link
Copy Markdown
Contributor

Part of #19519

This PR adds support for the FFT and complex operator family in the Relax TFLite frontend.

Key implementations:

  • Registered REAL, IMAG, COMPLEX_ABS, and RFFT2D to the TFLite op map.
  • Implemented convert_real and convert_imag which extract the real and imaginary parts of a complex tensor via strided_slice + squeeze along the last axis.
  • Implemented convert_complex_abs which computes sqrt(re^2 + im^2) using elementwise Relax ops.
  • All three ops adopt a unified representation convention: TFLite complex64 tensors (which have no native Relax dtype equivalent) are represented as float32[..., 2], where the last axis holds (real, imaginary) interleaved.

What is missing:

  • RFFT2D raises OpNotImplemented in this PR. A native
    relax.op.signal.rfft2d op with a C++ registered backend is required.
    topi.signal.dft exists as pure Python TE but has no
    TVM_REGISTER_GLOBAL entry and cannot be called via call_dps_packed.

Testing:

  • Added structural equality tests for REAL, IMAG, and COMPLEX_ABS in test_frontend_tflite.py following the verify(TestClass, Expected) pattern.
  • RFFT2D test covers frontend conversion correctness.
python3 -m pytest tests/python/relax/test_frontend_tflite.py -k "test_real or test_imag or test_complex_abs"

@fnhirwa fnhirwa marked this pull request as ready for review June 14, 2026 13:56

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread tests/python/relax/test_frontend_tflite.py Outdated
Comment thread tests/python/relax/test_frontend_tflite.py Outdated
Comment thread tests/python/relax/test_frontend_tflite.py Outdated

@tlopex tlopex left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 remains float32, so REAL/IMAG slice 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:

  • RFFT2D is registered in convert_map, but convert_rfft2d() always raises OpNotImplemented. The PR title/body says RFFT2D support and mentions an RFFT2D test, but the patch only adds tests for REAL, IMAG, and COMPLEX_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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants