[Relax][Frontend][TFLite] Add UNIDIRECTIONAL_SEQUENCE_RNN converter#19601
Merged
tlopex merged 1 commit intoMay 27, 2026
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request adds support for the UNIDIRECTIONAL_SEQUENCE_RNN operator in the TFLite frontend for TVM Relax, including implementation of the conversion logic and unit tests for various configurations. Feedback from the review suggests improving the implementation by supporting non-zero initial hidden states, handling optional bias tensors, and relaxing shape constraints to allow for dynamic batch sizes. A duplicate assertion in the test suite was also identified for removal.
Implements convert_unidirectional_sequence_rnn in the Relax TFLite
frontend. The op (BuiltinOperator 35) executes a simple RNN cell over a
time sequence using Relax primitives (matmul / add / activation).
Cell formula:
h_t = act(x_t @ W^T + h_{t-1} @ Wr^T + b)
Design notes:
- Inputs: input[batch,time,input_size], input_weights[units,input_size],
recurrent_weights[units,units], bias[units], hidden_state[batch,units]
- time_major=True input is transposed to batch-major before unrolling
- Activations supported: NONE, RELU, RELU6, TANH, SIGMOID
- Quantised variant raises OpNotImplemented (guard already present)
- For time=1 the split is skipped; squeeze is applied directly
- Time steps are unrolled at graph-construction time and outputs
stacked along axis=1
Tests added (3):
- test_unidirectional_sequence_rnn_none_activation: structural_equal
check with identity weights / zero bias, NONE activation, time=1
- test_unidirectional_sequence_rnn_relu_activation: shape check with
random weights, RELU activation, time=3
- test_unidirectional_sequence_rnn_time_major: shape check with
time_major=True input layout
Closes part of apache#19519
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
LudovicoYIN
force-pushed
the
relax/tflite/unidirectional-sequence-rnn
branch
from
0c021bb to
8906b06
Compare
Member
|
Hi @LudovicoYIN Could you have a look at the CI error? |
Contributor
Author
|
@tvm-bot rerun |
Contributor
|
Failed to re-run CI in https://github.com/apache/tvm/actions/runs/26486541101 Detailswith response |
Contributor
Author
|
Thanks your review. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds Relax TFLite frontend support for
UNIDIRECTIONAL_SEQUENCE_RNN(BuiltinOperator 35), claimed in #19519 Group A.The op executes a simple RNN cell over a time sequence. The converter unrolls the time steps at graph-construction time using Relax primitives.
Cell equation:
Changes
convert_unidirectional_sequence_rnnregistered inconvert_map(alphabetical, U-region afterUNPACK)input [batch, time, input_size],input_weights [num_units, input_size],recurrent_weights [num_units, num_units],bias [num_units],hidden_state [batch, num_units](variable, zero-initialised)[batch, time, num_units](always batch-major)convert_fused_activation_function)OpNotImplemented(not yet supported)Testing
Modern TF/Keras (2.x, Keras 3) no longer emits
UNIDIRECTIONAL_SEQUENCE_RNN;SimpleRNNwithunroll=Falselowers toWHILE+TensorList ops, andunroll=Trueexpands to elementwise ops. Tests therefore follow the same flatbuffer-construction pattern used by the StableHLO op PRs (#19536, #19587).Three tests added to
tests/python/relax/test_frontend_tflite.py:test_unidirectional_sequence_rnn_none_activation—tvm.ir.assert_structural_equalwith identity weights / zero bias, NONE activation, time=1test_unidirectional_sequence_rnn_relu_activation— shape check, random weights, RELU activation, time=3test_unidirectional_sequence_rnn_time_major— shape check,time_major=Trueinput layoutAll 3 tests pass. pre-commit (ASF header, ruff check, ruff format) all pass.
References