[Relax][Frontend][TFLite] Add EMBEDDING_LOOKUP_SPARSE converter#19652
Open
LudovicoYIN wants to merge 3 commits into
Open
[Relax][Frontend][TFLite] Add EMBEDDING_LOOKUP_SPARSE converter#19652LudovicoYIN wants to merge 3 commits into
LudovicoYIN wants to merge 3 commits into
Conversation
Add Relax TFLite frontend support for EMBEDDING_LOOKUP_SPARSE. This change implements the converter for the SUM, MEAN, and SQRTN combiners and supports higher-rank sparse indices. The converter lowers the sparse aggregation through scatter_nd to match TFLite operator semantics for the supported cases. The patch also adds handcrafted TFLite frontend tests covering the SUM, MEAN, and SQRTN combiners, along with a 3D indices case.
![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 implements support for the EMBEDDING_LOOKUP_SPARSE operator in the TVM Relax TFLite frontend, supporting SUM, MEAN, and SQRTN combiner types. It also adds comprehensive unit tests to verify the correctness of these combiners. Feedback on the implementation suggests optimizing the calculation of the safe denominator by using a scalar constant instead of creating a full-shaped constant tensor, leveraging the broadcasting capabilities of relax.op.maximum to improve efficiency and reduce compiled model size.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Use a scalar float32 constant when clamping the EMBEDDING_LOOKUP_SPARSE denominator. This keeps the Relax expression smaller by relying on broadcasting in maximum instead of materializing a full-shaped constant tensor.
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
1 participant
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
Add Relax TFLite frontend support for
EMBEDDING_LOOKUP_SPARSE.This PR adds a converter for
EMBEDDING_LOOKUP_SPARSEin the RelaxTFLite frontend. The implementation supports the
SUM,MEAN, andSQRTNcombiners and handles higher-rank sparse indices. The sparseaggregation is lowered through
scatter_ndto match TFLite operatorsemantics for the supported cases.
The PR also adds handcrafted TFLite frontend tests covering:
SUMMEANSQRTNTesting
Ran
tests/python/relax/test_frontend_tflite.py -k 'embedding_lookup_sparse'.Part of #19519