This is an automated email from the ASF dual-hosted git repository.
mshr-h pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 1fa4b8627d [Relax][ONNX] Add ReverseSequence operator support (#20924)
1fa4b8627d is described below
commit 1fa4b8627de240c59a2531bee7fb6a6f88981f59
Author: Anirudh Mittal <[email protected]>
AuthorDate: Sat Nov 7 23:31:25 2026 +0530
[Relax][ONNX] Add ReverseSequence operator support (#20924)
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 16 +++
tests/python/relax/test_frontend_onnx.py | 118 +++++++++++++++++++++++
tests/python/relax/test_frontend_onnx_backend.py | 1 +
3 files changed, 135 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 2f97ef839a..badb3f9975 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -3466,6 +3466,21 @@ class Tile(OnnxOpConverter):
return bb.emit_te(topi.dyn_tile, data, output_shape_vars, reps_len)
+class ReverseSequence(OnnxOpConverter):
+ """Converts an onnx ReverseSequence node into an equivalent Relax
expression."""
+
+ @classmethod
+ def _impl_v10(cls, bb, inputs, attr, params):
+ data = inputs[0]
+ seq_lengths = inputs[1]
+ batch_axis = int(attr.get("batch_axis", 1))
+ seq_axis = int(attr.get("time_axis", 0))
+
+ return relax.op.reverse_sequence(
+ data, seq_lengths, seq_axis=seq_axis, batch_axis=batch_axis
+ )
+
+
class Expand(OnnxOpConverter):
"""Converts an onnx Expand node into an equivalent Relax expression."""
@@ -6256,6 +6271,7 @@ def _get_convert_map():
"Pad": Pad,
"Split": Split,
"Tile": Tile,
+ "ReverseSequence": ReverseSequence,
"AveragePool": AveragePool,
"MaxPool": MaxPool,
"LpPool": LpPool,
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index d34d208b00..5e8c408d61 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -9386,6 +9386,124 @@ def test_tile_dynamic_repeats():
)
+def test_reverse_sequence_ir():
+ """Verify lowered Relax IR for the ReverseSequence ONNX operator."""
+
+ def make_reverse_sequence_model(x_shape, seq_lens_shape, time_axis=None,
batch_axis=None):
+ attrs = {}
+ if time_axis is not None:
+ attrs["time_axis"] = time_axis
+ if batch_axis is not None:
+ attrs["batch_axis"] = batch_axis
+ node = helper.make_node(
+ "ReverseSequence",
+ inputs=["x", "sequence_lens"],
+ outputs=["y"],
+ **attrs,
+ )
+ graph = helper.make_graph(
+ [node],
+ "reverse_sequence_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape),
+ helper.make_tensor_value_info("sequence_lens",
TensorProto.INT64, seq_lens_shape),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
x_shape)],
+ )
+ return helper.make_model(graph, producer_name="reverse_sequence_test")
+
+ # Default time_axis=0, batch_axis=1
+ model = make_reverse_sequence_model([4, 3], [3])
+ tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)
+
+ @I.ir_module
+ class ExpectedDefault:
+ @R.function
+ def main(
+ x: R.Tensor((4, 3), dtype="float32"),
+ sequence_lens: R.Tensor((3,), dtype="int64"),
+ ) -> R.Tensor((4, 3), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv: R.Tensor((4, 3), dtype="float32") = R.reverse_sequence(
+ x, sequence_lens, seq_axis=0, batch_axis=1
+ )
+ R.output(gv)
+ return gv
+
+ tvm_tirx.ir.assert_structural_equal(tvm_model, ExpectedDefault)
+
+ # Custom time_axis=1, batch_axis=0
+ model = make_reverse_sequence_model([2, 5], [2], time_axis=1, batch_axis=0)
+ tvm_model = from_onnx(model, opset=10, keep_params_in_input=True)
+
+ @I.ir_module
+ class ExpectedCustom:
+ @R.function
+ def main(
+ x: R.Tensor((2, 5), dtype="float32"),
+ sequence_lens: R.Tensor((2,), dtype="int64"),
+ ) -> R.Tensor((2, 5), dtype="float32"):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv: R.Tensor((2, 5), dtype="float32") = R.reverse_sequence(
+ x, sequence_lens, seq_axis=1, batch_axis=0
+ )
+ R.output(gv)
+ return gv
+
+ tvm_tirx.ir.assert_structural_equal(tvm_model, ExpectedCustom)
+
+
+def test_reverse_sequence():
+ """Verify numerical correctness of ReverseSequence against ONNX Runtime."""
+
+ def verify_reverse_sequence(x, sequence_lens, time_axis=None,
batch_axis=None, opset=10):
+ attrs = {}
+ if time_axis is not None:
+ attrs["time_axis"] = time_axis
+ if batch_axis is not None:
+ attrs["batch_axis"] = batch_axis
+ node = helper.make_node(
+ "ReverseSequence",
+ inputs=["x", "sequence_lens"],
+ outputs=["y"],
+ **attrs,
+ )
+ seq_lens_proto = TensorProto.INT64 if sequence_lens.dtype == np.int64
else TensorProto.INT32
+ graph = helper.make_graph(
+ [node],
+ "reverse_sequence_test",
+ inputs=[
+ helper.make_tensor_value_info("x", TensorProto.FLOAT,
list(x.shape)),
+ helper.make_tensor_value_info(
+ "sequence_lens", seq_lens_proto, list(sequence_lens.shape)
+ ),
+ ],
+ outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT,
list(x.shape))],
+ )
+ model = helper.make_model(graph, producer_name="reverse_sequence_test")
+ inputs = {"x": x, "sequence_lens": sequence_lens}
+ check_correctness(model, inputs=inputs, opset=opset)
+
+ # Example 1: 2D, time_axis=0, batch_axis=1
+ x = np.arange(16, dtype=np.float32).reshape(4, 4)
+ seq_lens = np.array([4, 3, 2, 1], dtype=np.int64)
+ verify_reverse_sequence(x, seq_lens, time_axis=0, batch_axis=1)
+
+ # Example 2: 2D, time_axis=1, batch_axis=0
+ seq_lens = np.array([1, 2, 3, 4], dtype=np.int64)
+ verify_reverse_sequence(x, seq_lens, time_axis=1, batch_axis=0)
+
+ # Default attributes (time_axis=0, batch_axis=1)
+ verify_reverse_sequence(x, np.array([4, 2, 1, 3], dtype=np.int64))
+
+ # 3D tensor
+ x_3d = np.arange(24, dtype=np.float32).reshape(2, 3, 4)
+ seq_lens_3d = np.array([3, 1], dtype=np.int64)
+ verify_reverse_sequence(x_3d, seq_lens_3d, time_axis=1, batch_axis=0)
+
+
def _generate_roi_cases():
return [
pytest.param(False, None, False, id="no_roi"),
diff --git a/tests/python/relax/test_frontend_onnx_backend.py
b/tests/python/relax/test_frontend_onnx_backend.py
index 6fa706c150..11ed7e1551 100644
--- a/tests/python/relax/test_frontend_onnx_backend.py
+++ b/tests/python/relax/test_frontend_onnx_backend.py
@@ -182,6 +182,7 @@ _INCLUDE_OPS = [
"or",
"reciprocal",
"relu",
+ "reversesequence",
"round",
"scatternd",
"sigmoid",
