[Relax][Frontend][TFLite] Add missing TFLite operator mappings by Aharrypotter · Pull Request #19813 · apache/tvm

Aharrypotter · 2026-06-17T09:38:56Z

Summary

Adds Relax TFLite frontend coverage for a batch of builtin operators tracked in
#19412. All lowerings reuse existing Relax ops; the change is limited to the
TFLite frontend converter and its tests.

Changes

Grouped by the #19412 checklist section:

Section A — drop-in operator mappings

SIGN -> relax.op.sign
BITWISE_XOR -> relax.op.bitwise_xor
RIGHT_SHIFT -> relax.op.right_shift
RELU_0_TO_1 -> relax.op.clip
BUCKETIZE -> relax.op.bucketize
UNIQUE -> relax.op.unique (unique values + inverse indices)

Section B — compose with existing Relax ops

RANK -> scalar rank constant for statically-ranked inputs
UNSORTED_SEGMENT_SUM -> extends the existing scatter-based segment lowering
(scatter_nd, reduction="add")
UNSORTED_SEGMENT_MAX -> same, reduction="max"

Section C — fix partial implementation

FAKE_QUANT: the handler called the non-existent relax.op.const and raised
AttributeError on every invocation; switch to relax.const so the op works.
(The builtin op carries only scalar min/max, so there is no per-channel/vector
min/max path.)

Tests

Adds coverage in tests/python/relax/test_frontend_tflite.py for the new
mappings and the FAKE_QUANT narrow-range vector path. Most tests use
structural-equal Relax IR checks; data-dependent and numeric regression cases
add focused output checks for UNIQUE and FAKE_QUANT.

python -m ruff format python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py
python -m ruff check  python/tvm/relax/frontend/tflite/tflite_frontend.py tests/python/relax/test_frontend_tflite.py
python -m pytest tests/python/relax/test_frontend_tflite.py -k "unique or sign or bitwise_xor or right_shift or bucketize or relu_0_to_1 or rank or unsorted_segment_sum or unsorted_segment_max or fake_quant"
python -m pytest tests/python/relax/test_frontend_tflite.py

Result:

ruff format: 2 files left unchanged
ruff check: All checks passed
target pytest: 22 passed, 523 deselected
full TFLite pytest: 545 passed

References

Issue [Tracking Issue][TFlite] Expand Missing / Partial Operator Coverage #19412: TFLite missing/partial operator coverage (sections A, B, C)

gemini-code-assist

Code Review

This pull request addresses several Relax TFLite frontend gaps by implementing mappings for operators such as SIGN, BITWISE_XOR, RIGHT_SHIFT, BUCKETIZE, RELU_0_TO_1, UNIQUE, RANK, UNSORTED_SEGMENT_SUM, and UNSORTED_SEGMENT_MAX, along with fixing FAKE_QUANT narrow-range vector conversion. Feedback suggests simplifying the RANK conversion by passing the Python integer directly to relax.const rather than wrapping it in a NumPy array, and using positional arguments instead of keyword arguments in relax.op.clip for consistency and compatibility.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-17T09:40:38Z

+        output_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())
+
+        rank = len(self.get_tensor_shape(input_tensors[0]))
+        return relax.const(np.array(rank, dtype=output_dtype), dtype=output_dtype)


We can simplify this by passing the Python integer rank directly to relax.const instead of wrapping it in a np.array. relax.const natively supports Python scalars with a specified dtype.

Suggested change

return relax.const(np.array(rank, dtype=output_dtype), dtype=output_dtype)

return relax.const(rank, dtype=output_dtype)

Addressed by returning relax.const(rank, dtype=output_dtype).

gemini-code-assist · 2026-06-17T09:40:38Z

+        if input_tensor.qnn_params:
+            in_f32 = self.dequantize(in_expr, input_tensor)
+            out = relax.op.clip(in_f32, min=0, max=1)
+            out = self.quantize(out, output_tensor)
+        else:
+            out = relax.op.clip(in_expr, min=0, max=1)


For consistency with other uses of relax.op.clip in this file (e.g., in convert_fake_quant), and to ensure compatibility with TVM's FFI wrappers, it is recommended to use positional arguments instead of keyword arguments (min and max).

Suggested change

if input_tensor.qnn_params:

in_f32 = self.dequantize(in_expr, input_tensor)

out = relax.op.clip(in_f32, min=0, max=1)

out = self.quantize(out, output_tensor)

else:

out = relax.op.clip(in_expr, min=0, max=1)

if input_tensor.qnn_params:

in_f32 = self.dequantize(in_expr, input_tensor)

out = relax.op.clip(in_f32, 0, 1)

out = self.quantize(out, output_tensor)

else:

out = relax.op.clip(in_expr, 0, 1)

Addressed by switching the new relax.op.clip calls to positional arguments.

tlopex · 2026-06-17T18:34:25Z

Could you resolve the conflict so that we can merge it in?

Aharrypotter · 2026-06-18T02:48:57Z

Rebased the PR onto the latest main and resolved the branch conflict in the
TFLite converter map.

Also refreshed the PR description with the post-rebase validation results.

Aharrypotter · 2026-06-23T02:54:28Z

Updated the PR description for clearer scope: the changes are now grouped by the
#19412 checklist section — Section A (drop-in mappings: SIGN, BITWISE_XOR,
RIGHT_SHIFT, RELU_0_TO_1, BUCKETIZE, UNIQUE), Section B (RANK,
UNSORTED_SEGMENT_SUM/MAX), and Section C (the FAKE_QUANT crash fix:
relax.op.const → relax.const) —
with a link back to the tracking issue. No code change; this only makes it
clearer which tracked items this PR covers.

## Summary This PR adds Relax TFLite frontend support for dynamic (runtime) scalar bounds in the `RANGE` operator, addressing the `RANGE` "fix partial implementations" item from #19412 section C. `convert_range` previously lowered only **constant** `start`, `limit`, and `delta` to `relax.op.arange` and raised `OpNotImplemented` for runtime scalar bounds (the guard added in #19401). Models that compute RANGE bounds at runtime could therefore not be imported. This PR makes the dynamic path work for both integer and float bounds, ascending or descending, without adding a new Relax op. The change is limited to the `RANGE` converter and its test. #19813 added a batch of missing TFLite operator mappings but did not touch this partial-implementation item; this PR closes it. ## Design ### Dynamic scalar bounds via count-lift `relax.op.arange` only accepts compile-time `PrimExpr` bounds. The frontend already has a runtime-scalar -> symbolic-dimension bridge (`relax.op.tensor_to_shape` + `match_cast`, as used by `_get_shape_expr_from_tensor`), so no new op is needed. Rather than feed symbolic bounds straight into `arange`, the converter computes the element **count** in-graph and lifts that single value to one symbolic output dimension `L`, then rebuilds the values as `arange(0, L) * delta + start`. Lifting the count (instead of the bounds) keeps the declared and runtime output lengths equal by construction: `arange`'s struct-info length formula (`InferTypeArange`) has no negative-step branch, so feeding symbolic bounds directly would mis-declare descending ranges relative to the TOPI runtime length. The count is `max(0, ceil((limit - start) / delta))`, computed per dtype: - **integer**: `-floor_divide(start - limit, delta)` — exact, sign-agnostic, and free of float-precision loss; equal to `ceil((limit - start) / delta)`. - **float**: `ceil((limit - start) / delta)`. Constant (all-bounds-constant) RANGE keeps the existing direct-`arange` path unchanged. ## Operator Support | Operator | TFLite inputs | Relax lowering | Supported subset | |---|---|---|---| | `RANGE` | scalar `start`, `limit`, `delta` | `relax.op.arange` (constant bounds); count-lift + `arange(0, L) * delta + start` (dynamic bounds) | int and float, constant or runtime scalar bounds, ascending or descending | ## Tests The dynamic test compiles the imported module and runs it on the Relax VM, comparing the output against `numpy.arange`. The constant-bound structural test is unchanged. | Test | Coverage | |---|---| | `test_range` | constant scalar bounds (existing, unchanged) | | `test_range_dynamic_scalar_inputs` | runtime scalar bounds: int and float, ascending and descending | Local validation: ```bash python -m ruff format --check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m ruff check \ python/tvm/relax/frontend/tflite/tflite_frontend.py \ tests/python/relax/test_frontend_tflite.py python -m pytest \ tests/python/relax/test_frontend_tflite.py -k range -q python -m pytest \ tests/python/relax/test_frontend_tflite.py -q ``` Result: ```text ruff format --check: 2 files already formatted ruff check: All checks passed range tests: 12 passed, 536 deselected full TFLite pytest: 548 passed ``` ## References - Issue #19412 section C: fix partial TFLite operator implementations (`RANGE`) - PR #19401: added the `RANGE` dynamic-scalar guard and its test - PR #18868: introduced the Relax TFLite frontend and `convert_range`

gemini-code-assist Bot reviewed Jun 17, 2026

View reviewed changes

Aharrypotter force-pushed the codex/tflite-operator-gap-closure branch from 5843680 to 864935f Compare June 17, 2026 09:49

Aharrypotter changed the title ~~[Relax][Frontend][TFLite] Close remaining TFLite operator gaps~~ [Relax][Frontend][TFLite] Add missing TFLite operator mappings Jun 17, 2026

Aharrypotter marked this pull request as ready for review June 17, 2026 10:41

Aharrypotter marked this pull request as draft June 17, 2026 10:47

Aharrypotter marked this pull request as ready for review June 17, 2026 11:01

tlopex approved these changes Jun 17, 2026

View reviewed changes

Aharrypotter added 2 commits June 18, 2026 10:25

[Relax][Frontend][TFLite] Add remaining operator mappings

1fecb08

[Relax][Frontend][TFLite] Apply ruff format

0ff302e

Aharrypotter force-pushed the codex/tflite-operator-gap-closure branch from c786073 to 0ff302e Compare June 18, 2026 02:36

Aharrypotter requested a review from tlopex June 18, 2026 02:59

tlopex merged commit 0bc3a6f into apache:main Jun 18, 2026
9 of 10 checks passed

Aharrypotter mentioned this pull request Jun 23, 2026

[Relax][Frontend][TFLite] Support dynamic RANGE scalar bounds #19867

Merged

Aharrypotter mentioned this pull request Jun 24, 2026

[Tracking Issue][TFlite] Expand Missing / Partial Operator Coverage #19412

Open

33 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Relax][Frontend][TFLite] Add missing TFLite operator mappings#19813

[Relax][Frontend][TFLite] Add missing TFLite operator mappings#19813
tlopex merged 2 commits into
apache:mainfrom
Aharrypotter:codex/tflite-operator-gap-closure

Aharrypotter commented Jun 17, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 17, 2026

Uh oh!

Aharrypotter Jun 17, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot Jun 17, 2026

Uh oh!

Aharrypotter Jun 17, 2026 •

edited

Loading

Uh oh!

tlopex commented Jun 17, 2026

Uh oh!

Aharrypotter commented Jun 18, 2026

Uh oh!

Uh oh!

Aharrypotter commented Jun 23, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	return relax.const(np.array(rank, dtype=output_dtype), dtype=output_dtype)
	return relax.const(rank, dtype=output_dtype)

Uh oh!

Conversation

Aharrypotter commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Tests

References

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

Aharrypotter Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

Aharrypotter Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tlopex commented Jun 17, 2026

Uh oh!

Aharrypotter commented Jun 18, 2026

Uh oh!

Uh oh!

Aharrypotter commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Aharrypotter commented Jun 17, 2026 •

edited

Loading

Aharrypotter Jun 17, 2026 •

edited

Loading

Aharrypotter Jun 17, 2026 •

edited

Loading

Aharrypotter commented Jun 23, 2026 •

edited

Loading