Skip to content

[Relax][Frontend][TFLite] Add tests coverage for SPACE_TO_BATCH_ND and BATCH_TO_SPACE_ND#19499

Merged
tlopex merged 2 commits into
apache:mainfrom
rknastenka:tests
May 4, 2026
Merged

[Relax][Frontend][TFLite] Add tests coverage for SPACE_TO_BATCH_ND and BATCH_TO_SPACE_ND#19499
tlopex merged 2 commits into
apache:mainfrom
rknastenka:tests

Conversation

@rknastenka

@rknastenka rknastenka commented May 3, 2026

Copy link
Copy Markdown
Contributor

Changes
Add tests in test_frontend_tflite.py.
Lower SPACE_TO_BATCH_ND / BATCH_TO_SPACE_ND through TOPI in tflite_frontend.py.
Use tf.raw_ops.BatchToSpaceND in the test because tf.batch_to_space_nd is not available in this TF build.

Why the TFLite frontend changed
The frontend was calling relax.op.nn.space_to_batch_nd / relax.op.nn.batch_to_space_nd, which aren’t implemented in this checkout. I updated the TFLite frontend to lower these ops via TOPI packed calls so conversion works and the new tests can pass.

Test:

pytest test_frontend_tflite.py -k "test_space_to_batch_nd or test_batch_to_space_nd"

related to #18971

@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 to implement batch_to_space_nd and space_to_batch_nd using call_dps_packed with TOPI functions, and adds unit tests for these operations. Feedback includes a correction for an incorrect argument count in the batch_to_space_nd call and a recommendation to use to_int_list for shape and padding parameters to ensure type safety within the Relax IR.

Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py Outdated
Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.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.

LGTM

@tlopex tlopex merged commit 87bf302 into apache:main May 4, 2026
10 checks passed
@rknastenka rknastenka deleted the tests branch May 4, 2026 06:08
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