[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

Author: Aharrypotter

include/tvm/relax/attrs/manipulate.h

@@ -195,6 +195,23 @@ struct FlipAttrs : public AttrsNode {
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.FlipAttrs", FlipAttrs, AttrsNode);
 };  // struct FlipAttrs
 
+/*! \brief Attributes used in reverse_sequence operators */
+struct ReverseSequenceAttrs : public AttrsNode {
+  int64_t seq_axis;
+  int64_t batch_axis;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<ReverseSequenceAttrs>()
+        .def_ro("seq_axis", &ReverseSequenceAttrs::seq_axis,
+                "The axis along which to reverse variable length slices.")
+        .def_ro("batch_axis", &ReverseSequenceAttrs::batch_axis,
+                "The axis that indexes the batch.");
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.attrs.ReverseSequenceAttrs", ReverseSequenceAttrs,
+                                    AttrsNode);
+};  // struct ReverseSequenceAttrs
+
 /*! \brief Attributes used in gather_elements operators */
 struct GatherElementsAttrs : public AttrsNode {
   int64_t axis;

python/tvm/relax/frontend/tflite/tflite_frontend.py

@@ -5614,10 +5614,10 @@ def convert_unpack(self, op):
         else:
             splitted = relax.op.split(in_expr, indices_or_sections=num_unpacks, axis=unpack_axis)
             squeezed = relax.Tuple(
-                relax.Tuple(
-                    [_op.squeeze(split_item, axis=squeeze_axis) for split_item in splitted]
-                ),
-                len(splitted),
+                [
+                    _op.squeeze(relax.TupleGetItem(splitted, i), axis=squeeze_axis)
+                    for i in range(num_unpacks)
+                ]
             )
 
         return squeezed

python/tvm/relax/op/__init__.py

@@ -108,6 +108,7 @@
     permute_dims,
     repeat,
     reshape,
+    reverse_sequence,
     scatter_elements,
     scatter_nd,
     slice_scatter,

python/tvm/relax/op/manipulate.py

@@ -460,6 +460,31 @@ def flip(data, axis):
     return _ffi_api.flip(data, axis)  # type: ignore
 
 
+def reverse_sequence(data: Expr, seq_lengths: Expr, seq_axis: int = 1, batch_axis: int = 0) -> Expr:
+    """Reverses variable length slices.
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input tensor.
+
+    seq_lengths : relax.Expr
+        A 1-D tensor containing sequence lengths for each batch.
+
+    seq_axis : int
+        The axis along which to reverse variable length slices.
+
+    batch_axis : int
+        The axis that indexes the batch.
+
+    Returns
+    -------
+    ret : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.reverse_sequence(data, seq_lengths, seq_axis, batch_axis)  # type: ignore
+
+
 def gather_elements(data: Expr, indices: Expr, axis: int = 0) -> Expr:
     """Gather elements from data according to indices along the specified axis.
 

python/tvm/relax/op/op_attrs.py

@@ -206,6 +206,11 @@ class FlipAttrs(Attrs):
     """Attributes for flip operator"""
 
 
+@tvm_ffi.register_object("relax.attrs.ReverseSequenceAttrs")
+class ReverseSequenceAttrs(Attrs):
+    """Attributes for reverse_sequence operator"""
+
+
 @tvm_ffi.register_object("relax.attrs.PadAttrs")
 class PadAttrs(Attrs):
     """Attributes used in pad operator"""

python/tvm/relax/script/builder/ir.py

@@ -153,6 +153,7 @@
     quantize,
     repeat,
     reshape,
+    reverse_sequence,
     right_shift,
     round,
     rsqrt,
@@ -941,6 +942,7 @@ def dtype(value: py_str | DataType) -> Expr:
     "quantize",
     "repeat",
     "reshape",
+    "reverse_sequence",
     "rewriter",
     "right_shift",
     "rocm",

python/tvm/relax/transform/legalize_ops/manipulate.py

@@ -170,6 +170,18 @@ def _flip(bb: BlockBuilder, call: Call) -> Expr:
     return bb.call_te(topi.flip, call.args[0], int(call.attrs.axis))
 
 
+@register_legalize("relax.reverse_sequence")
+def _reverse_sequence(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.reverse_sequence,
+        call.args[0],
+        call.args[1],
+        int(call.attrs.seq_axis),
+        int(call.attrs.batch_axis),
+        primfunc_name_hint="reverse_sequence",
+    )
+
+
 @register_legalize("relax.gather_elements")
 def _gather_elements(bb: BlockBuilder, call: Call) -> Expr:
     return bb.call_te(topi.gather, call.args[0], int(call.attrs.axis), call.args[1])

src/relax/op/tensor/manipulate.cc

@@ -51,6 +51,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
   RepeatAttrs::RegisterReflection();
   TileAttrs::RegisterReflection();
   FlipAttrs::RegisterReflection();
+  ReverseSequenceAttrs::RegisterReflection();
   GatherElementsAttrs::RegisterReflection();
   GatherNDAttrs::RegisterReflection();
   IndexPutAttrs::RegisterReflection();
@@ -2071,6 +2072,96 @@ TVM_REGISTER_OP("relax.flip")
     .set_attr<FRelaxInferLayout>("FRelaxInferLayout", InferLayoutFlip)
     .set_attr<bool>("FPurity", true);
 
+/* relax.reverse_sequence */
+
+Expr reverse_sequence(Expr data, Expr seq_lengths, int64_t seq_axis, int64_t batch_axis) {
+  auto attrs = ffi::make_object<ReverseSequenceAttrs>();
+  attrs->seq_axis = seq_axis;
+  attrs->batch_axis = batch_axis;
+  static const Op& op = Op::Get("relax.reverse_sequence");
+  return Call(op, {std::move(data), std::move(seq_lengths)}, Attrs{attrs}, {});
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+  namespace refl = tvm::ffi::reflection;
+  refl::GlobalDef().def("relax.op.reverse_sequence", reverse_sequence);
+}
+
+Type InferTypeReverseSequence(const Call& call, const BlockBuilder& ctx) {
+  if (call->args.size() != 2) {
+    TVM_FFI_VISIT_THROW(ValueError, call) << "ReverseSequence op should take 2 arguments";
+  }
+  TensorType data_ty = GetInputTensorType(call, 0, ctx);
+  TensorType seq_lengths_ty = GetInputTensorType(call, 1, ctx);
+
+  if (!seq_lengths_ty->IsUnknownNdim() && seq_lengths_ty->ndim != 1) {
+    TVM_FFI_VISIT_THROW(ValueError, call)
+        << "ReverseSequence requires seq_lengths to be 1-D. However, seq_lengths has ndim "
+        << seq_lengths_ty->ndim;
+  }
+  if (!seq_lengths_ty->dtype.is_void() && !seq_lengths_ty->dtype.is_int()) {
+    TVM_FFI_VISIT_THROW(ValueError, call)
+        << "ReverseSequence requires seq_lengths to have dtype int32 or int64. However, "
+           "seq_lengths has dtype "
+        << seq_lengths_ty->dtype;
+  }
+  if (seq_lengths_ty->dtype.is_int() && seq_lengths_ty->dtype.bits() != 32 &&
+      seq_lengths_ty->dtype.bits() != 64) {
+    TVM_FFI_VISIT_THROW(ValueError, call)
+        << "ReverseSequence requires seq_lengths to have dtype int32 or int64. However, "
+           "seq_lengths has dtype "
+        << seq_lengths_ty->dtype;
+  }
+
+  const auto* attrs = call->attrs.as<ReverseSequenceAttrs>();
+  int64_t seq_axis = attrs->seq_axis;
+  int64_t batch_axis = attrs->batch_axis;
+  if (!data_ty->IsUnknownNdim()) {
+    int ndim = data_ty->ndim;
+    auto check_axis = [&](int64_t axis, ffi::String axis_name) {
+      if (axis < -ndim || axis >= ndim) {
+        TVM_FFI_VISIT_THROW(ValueError, call)
+            << "ReverseSequence requires " << axis_name
+            << " to belong to range [-ndim, ndim). However, the axis is " << axis
+            << ", while ndim is " << ndim;
+      }
+    };
+    check_axis(seq_axis, "seq_axis");
+    check_axis(batch_axis, "batch_axis");
+
+    if (batch_axis < 0) {
+      batch_axis += ndim;
+    }
+
+    if (data_ty->shape.defined() && seq_lengths_ty->shape.defined()) {
+      const auto* data_shape_ty = GetTypeAs<ShapeTypeNode>(data_ty->shape.value());
+      const auto* seq_lengths_shape_ty = GetTypeAs<ShapeTypeNode>(seq_lengths_ty->shape.value());
+      if (data_shape_ty != nullptr && seq_lengths_shape_ty != nullptr &&
+          data_shape_ty->values.defined() && seq_lengths_shape_ty->values.defined()) {
+        PrimExpr batch_extent = data_shape_ty->values.value()[batch_axis];
+        PrimExpr seq_lengths_extent = seq_lengths_shape_ty->values.value()[0];
+        if (ctx->GetAnalyzer()->CanProve(seq_lengths_extent != batch_extent)) {
+          TVM_FFI_VISIT_THROW(ValueError, call)
+              << "ReverseSequence requires seq_lengths.shape[0] to equal the batch axis extent. "
+                 "However, seq_lengths.shape[0] is "
+              << seq_lengths_extent << ", while data.shape[" << batch_axis << "] is "
+              << batch_extent;
+        }
+      }
+    }
+  }
+
+  return data_ty;
+}
+
+TVM_REGISTER_OP("relax.reverse_sequence")
+    .set_attrs_type<ReverseSequenceAttrs>()
+    .set_num_inputs(2)
+    .add_argument("data", "Tensor", "The input tensor.")
+    .add_argument("seq_lengths", "Tensor", "The sequence length tensor.")
+    .set_attr<FInferType>("FInferType", InferTypeReverseSequence)
+    .set_attr<bool>("FPurity", true);
+
 /* relax.gather_elements */
 
 Expr gather_elements(Expr data, Expr indices, int axis) {

src/relax/op/tensor/manipulate.h

@@ -181,6 +181,16 @@ Expr tile(Expr data, ffi::Array<int64_t> repeats);
  */
 Expr flip(Expr data, int64_t axis);
 
+/*!
+ * \brief Reverses variable length slices along seq_axis.
+ * \param data The input tensor.
+ * \param seq_lengths A 1-D tensor containing sequence lengths for each batch.
+ * \param seq_axis The axis along which to reverse.
+ * \param batch_axis The axis that indexes the batch.
+ * \return The computed result.
+ */
+Expr reverse_sequence(Expr data, Expr seq_lengths, int64_t seq_axis, int64_t batch_axis);
+
 /*!
  * \brief Gather elements from a tensor using indices.
  * \param data The input tensor.

src/topi/transform.cc

@@ -58,8 +58,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
                   })
       .def_packed("topi.reverse_sequence",
                   [](ffi::PackedArgs args, ffi::Any* rv) {
+                    int batch_axis = args.size() >= 4 ? args[3].cast<int>() : 0;
                     *rv = reverse_sequence(args[0].cast<te::Tensor>(), args[1].cast<te::Tensor>(),
-                                           args[2].cast<int>());
+                                           args[2].cast<int>(), batch_axis);
                   })
       .def_packed("topi.reshape",
                   [](ffi::PackedArgs args, ffi::Any* rv) {