Skip to content

[Relax][Frontend][TFLite] Add SCATTER_ND operator for Relax TFLite#19490

Merged
tlopex merged 2 commits intoapache:mainfrom
rknastenka:con2
May 1, 2026
Merged

[Relax][Frontend][TFLite] Add SCATTER_ND operator for Relax TFLite#19490
tlopex merged 2 commits intoapache:mainfrom
rknastenka:con2

Conversation

@rknastenka
Copy link
Copy Markdown
Contributor

This PR adds support for the SCATTER_ND operator in the Relax TFLite frontend.

Key Changes:

  • Added handler convert_scatter_nd to parse indices and updates.
  • Explicitly handles static vs dynamic shape tensor extraction via to_int_list and relax.op.tensor_to_shape.
  • Uses relax.op.zeros to initialize the base array based on the updates precision dtype.
  • Mapped SCATTER_ND to the corresponding relax.op.scatter_nd() target.
  • Registered the translator into convert_map and provided the matching unit test in test_frontend_tflite.py.

Testing:

Passed unit tests

pytest tests/python/relax/test_frontend_tflite.py::test_scatter_nd

Related to #19412

Copy link
Copy Markdown
Contributor

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

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 the SCATTER_ND operator to the TFLite frontend in Relax and includes a new test case to verify its functionality. The review feedback highlights two necessary corrections: the indices tensor needs to be permuted to align with Relax's expected layout (where the index depth is the first dimension), and the shape tensor must be explicitly cast to int64 when handled dynamically to satisfy the requirements of the tensor_to_shape operator.

Comment thread python/tvm/relax/frontend/tflite/tflite_frontend.py
@tlopex tlopex merged commit fbbbae9 into apache:main May 1, 2026
9 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