The GitHub Actions job "Teams" on tvm.git/main has succeeded. Run started by GitHub user tlopex (triggered by tlopex).
Head commit for run: 4b0a0397e2805fb1bea8dabaa0412463887859d3 / Hongyi Wu <[email protected]> [Relax][TFLite] Add remaining operator tests and reverse_sequence op (#19814) ## Summary This PR adds focused Relax TFLite frontend coverage for the remaining non-quantized builtin operators tracked by #18971: - `SQUEEZE` - `REVERSE_SEQUENCE` - `UNPACK` - `ZEROS_LIKE` The tests manually build minimal TFLite flatbuffers and compare the imported Relax IR with `tvm.ir.assert_structural_equal`. This keeps the coverage on the frontend importer itself, without depending on TensorFlow converter rewrites or constant folding. The PR also adds first-class Relax support for `reverse_sequence`. TFLite `REVERSE_SEQUENCE` was previously routed through: ```text R.call_dps_packed("topi.reverse_sequence", ...) ``` That is not executable as a runtime packed call because `topi.reverse_sequence` is a TE compute and expects TE tensors during lowering. The frontend now emits `R.reverse_sequence`, and `LegalizeOps` lowers it through TOPI to TIR: ```text TFLite REVERSE_SEQUENCE -> R.reverse_sequence -> LegalizeOps -> topi.reverse_sequence -> R.call_tir ``` ## Design ### TFLite Operator Tests The new TFLite tests use hand-built flatbuffers for small importer fixtures: - `SQUEEZE` checks axis handling and direct Relax `squeeze` lowering. - `REVERSE_SEQUENCE` checks import to `R.reverse_sequence`, rejects the old `R.call_dps_packed("topi.reverse_sequence", ...)` path, compiles the module, and runs it with the VM. - `UNPACK` checks multi-output lowering through Relax tuple output handling. - `ZEROS_LIKE` checks direct Relax zero-like tensor creation. ### Relax reverse_sequence Operator The PR adds a public Relax operator: ```python relax.op.reverse_sequence(data, seq_lengths, seq_axis=1, batch_axis=0) ``` The operator uses `ReverseSequenceAttrs` with `seq_axis` and `batch_axis`. Type inference preserves the input tensor's shape, dtype, and vdevice, and validates the statically known constraints: - `data` must be a tensor. - `seq_lengths` must be a 1-D tensor. - `seq_lengths` dtype must be `int32` or `int64`. - `seq_axis` and `batch_axis` must be in `[-ndim, ndim)` when the input rank is known. - `seq_lengths.shape[0]` must match the batch-axis extent when both shapes are statically available. The op is exported through Python as `relax.op.reverse_sequence` and through the script builder as `R.reverse_sequence`. ### Legalization `relax.reverse_sequence` is registered in `LegalizeOps` and lowered with `bb.call_te`: ```python bb.call_te( topi.reverse_sequence, data, seq_lengths, seq_axis, batch_axis, primfunc_name_hint="reverse_sequence", ) ``` This produces `R.call_tir` in the legalized Relax module, keeping runtime execution on the normal TOPI/TIR path. ### TOPI Packed Registration The Python TOPI wrapper already accepts `batch_axis`: ```python topi.reverse_sequence(a, seq_lengths, seq_axis=1, batch_axis=0) ``` The C++ packed registration only forwarded the first three arguments, so Python calls that provided `batch_axis` would drop it before reaching the TOPI compute. The registration now forwards the fourth argument and keeps the old three-argument call form compatible by defaulting `batch_axis=0`. ## Operator Support | Operator | TFLite options | Relax lowering | Supported subset | |---|---|---|---| | `SQUEEZE` | `SqueezeOptions.SqueezeDims()` | `R.squeeze` | static squeeze axes from TFLite options | | `REVERSE_SEQUENCE` | `ReverseSequenceOptions.SeqDim()`, `BatchDim()` | `R.reverse_sequence` legalized to TOPI/TIR | tensor input, 1-D int32/int64 `seq_lengths`, valid `seq_axis` and `batch_axis` | | `UNPACK` | `UnpackOptions.Axis()`, `Num()` | Relax tuple output | static axis and output count from TFLite options | | `ZEROS_LIKE` | none | `R.zeros_like` | tensor input | ## Not Included - Quantized TFLite `REVERSE_SEQUENCE` support. - A runtime DPS packed implementation for `topi.reverse_sequence`. - Changes to TOPI compute semantics. - ONNX `ReverseSequence` importer support. ## Tests The tests cover both the TFLite frontend fixtures and the new Relax op: | Test | Coverage | |---|---| | `test_squeeze` | imports TFLite `SQUEEZE` to Relax `squeeze` | | `test_reverse_sequence` | imports TFLite `REVERSE_SEQUENCE` to `R.reverse_sequence`, avoids the old TOPI DPS packed call, compiles, and runs through VM | | `test_unpack` | imports TFLite `UNPACK` as multi-output Relax tuple handling | | `test_zeros_like` | imports TFLite `ZEROS_LIKE` to Relax `zeros_like` | | `test_op_correctness` | `relax.op.reverse_sequence(...).op` resolves to `relax.reverse_sequence` | | `test_reverse_sequence_infer_ty` | static shape, unknown dtype, unknown ndim, symbolic shape, and vdevice propagation | | `test_reverse_sequence_infer_ty_wrong_inputs` | non-tensor `seq_lengths`, wrong rank, wrong dtype, invalid axes, and static batch mismatch | | `test_reverse_sequence` in `test_transform_legalize_ops_manipulate.py` | `LegalizeOps` emits `R.call_tir` and exercises `seq_axis=0, batch_axis=1` | Local validation: ```bash python -m pytest tests/python/relax/test_op_manipulate.py \ -k reverse_sequence -q python -m pytest tests/python/relax/test_transform_legalize_ops_manipulate.py \ -k reverse_sequence -q python -m pytest --noconftest tests/python/relax/test_frontend_tflite.py \ -k "reverse_sequence or squeeze or unpack or zeros_like" -q ``` Result: ```text cmake build: passed py_compile: passed ruff format --check: 9 files already formatted ruff check: All checks passed clang-format --dry-run --Werror: passed pre-commit run --files: passed test_op_manipulate.py -k reverse_sequence: 3 passed test_transform_legalize_ops_manipulate.py -k reverse_sequence: 1 passed test_frontend_tflite.py -k "reverse_sequence or squeeze or unpack or zeros_like": 4 passed ``` ## References - Issue #18971: TFLite non-quantized operator unit-test coverage - TFLite `REVERSE_SEQUENCE` builtin semantics Report URL: https://github.com/apache/tvm/actions/runs/28051154831 With regards, GitHub Actions via GitBox --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
