Skip to content

[Relax][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT#19373

Merged
tlopex merged 2 commits into
apache:mainfrom
swjng:test/tflite-padding-topk-onehot
Apr 9, 2026
Merged

[Relax][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT#19373
tlopex merged 2 commits into
apache:mainfrom
swjng:test/tflite-padding-topk-onehot

Conversation

@swjng

@swjng swjng commented Apr 9, 2026

Copy link
Copy Markdown
Contributor

Part of #18971

Two bugs in the TFLite Relax frontend converters are fixed, and unit tests
are added for the Padding / Sparse / Other category operators claimed in
that tracking issue.

Bug fixes

convert_mirror_pad

Called relax.op.nn.mirror_pad which does not exist in the Relax op
namespace. Replaced with relax.op.nn.pad using pad_mode="reflect" for
REFLECT mode (the modes are semantically equivalent). SYMMETRIC mode raises
OpAttributeUnImplemented as there is no direct Relax equivalent.

convert_one_hot

  • on_value and off_value were passed as Expr (constant tensor nodes),
    but relax.op.one_hot requires PrimValue arguments.
  • An extra dtype positional argument was passed, which the function
    signature does not accept.

Fixed by extracting the scalar from each tensor buffer and wrapping it in
relax.PrimValue with the correct dtype via tvm.tirx.FloatImm /
tvm.tirx.IntImm.

Tests added

Each test uses the verify() + tf.Module pattern and includes an explicit
expected IRModule verified with tvm.ir.assert_structural_equal.

Operator TFLite op Notes
test_pad PAD constant zero padding
test_pad_v2 PADV2 explicit constant_values=5.0
test_mirror_pad MIRROR_PAD REFLECT mode
test_topk_v2 TOPK_V2 returns top-k values
test_one_hot ONE_HOT float32 on/off values, depth=4

@swjng swjng changed the title [BugFix][Test][TFLite] Fix and test PAD/PADV2/MIRROR_PAD/TOPK_V2/ONE_HOT converters [BugFix][TFLite] Fix MIRROR_PAD and ONE_HOT converters and add unit tests Apr 9, 2026
@swjng swjng changed the title [BugFix][TFLite] Fix MIRROR_PAD and ONE_HOT converters and add unit tests [BugFix][TFLite] Fix MIRROR_PAD and ONE_HOT converters and add unit tests Apr 9, 2026
…HOT converters

Two bugs in the TFLite Relax frontend converters are fixed:

1. convert_mirror_pad called relax.op.nn.mirror_pad which does not exist
   in the Relax op namespace. Replace with relax.op.nn.pad using
   pad_mode="reflect" for REFLECT mode, and raise OpAttributeUnImplemented
   for SYMMETRIC mode (no equivalent in Relax).

2. convert_one_hot passed on_value and off_value as Expr (constant tensor
   nodes) where relax.op.one_hot requires PrimValue, and included an extra
   dtype positional argument that the function signature does not accept.
   Fix by extracting the scalar from the tensor buffer and wrapping it in
   relax.PrimValue with the correct dtype.

Unit tests are added for PAD, PADV2, MIRROR_PAD, TOPK_V2, and ONE_HOT
following the verify() + tf.Module pattern. Each test includes an explicit
expected IRModule verified with tvm.ir.assert_structural_equal.

Partially closes apache#18971.
@swjng swjng changed the title [BugFix][TFLite] Fix MIRROR_PAD and ONE_HOT converters and add unit tests [BugFix][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT Apr 9, 2026
@swjng swjng force-pushed the test/tflite-padding-topk-onehot branch from c42e8a9 to 44b64b8 Compare April 9, 2026 01:32
@swjng swjng changed the title [BugFix][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT [BugFix][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT Apr 9, 2026

@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 updates the TFLite frontend in TVM Relax by refactoring the MIRROR_PAD and ONE_HOT operators. MIRROR_PAD now utilizes relax.op.nn.pad for reflection padding, while ONE_HOT has been updated to wrap its on/off values as PrimValue. Additionally, several new tests were added to cover padding and one-hot operations. Review feedback suggests casting padding values to Python integers for better compatibility and simplifying the extraction of scalar values from numpy arrays in the one-hot conversion.

Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
@swjng swjng changed the title [BugFix][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT [Relax][TFLite] Fix MIRROR_PAD/ONE_HOT converters and add tests for PAD, PADV2, MIRROR_PAD, TOPK_V2, ONE_HOT Apr 9, 2026

@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.

LGTM Thanks!

@tlopex tlopex merged commit fb6453a into apache:main Apr 9, 2026
10 checks passed
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