From fed9e4fc79d1e7a22768938132eea512fb29d9c9 Mon Sep 17 00:00:00 2001 From: Anirudh-65 Date: Fri, 6 Nov 2026 01:10:01 +0530 Subject: [PATCH 1/2] [Relax][ONNX] Add ReverseSequence operator support This commit adds support for the ONNX ReverseSequence operator in the Relax ONNX frontend. Key changes: - In `python/tvm/relax/frontend/onnx/onnx_frontend.py`, added the `ReverseSequence` converter class (`_impl_v10`) mapping ONNX attributes `time_axis` (default 0) to Relax `seq_axis` and `batch_axis` (default 1) to Relax `batch_axis`. - Registered "ReverseSequence" in `_get_convert_map()`. - In `tests/python/relax/test_frontend_onnx.py`, added structural IR equivalence tests (`test_reverse_sequence_ir`) and numerical correctness tests against ONNX Runtime (`test_reverse_sequence`). - In `tests/python/relax/test_frontend_onnx_backend.py`, added "reversesequence" to `_INCLUDE_OPS` to enable official ONNX node tests. --- .../tvm/relax/frontend/onnx/onnx_frontend.py | 16 +++ tests/python/relax/test_frontend_onnx.py | 118 ++++++++++++++++++ .../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 8866019c5523..0ba0bc057816 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -3450,6 +3450,21 @@ def _impl_v13(cls, bb, inputs, attr, params): 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.""" @@ -6239,6 +6254,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 bba6ef4faf0e..8396db9e1822 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -9326,6 +9326,124 @@ def main( ) +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, 0, 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 6fa706c150ca..11ed7e155165 100644 --- a/tests/python/relax/test_frontend_onnx_backend.py +++ b/tests/python/relax/test_frontend_onnx_backend.py @@ -182,6 +182,7 @@ def supports_device(cls, device: str) -> bool: "or", "reciprocal", "relu", + "reversesequence", "round", "scatternd", "sigmoid", From 12d7cb4f2113c46530712fa320b22b485af3af5b Mon Sep 17 00:00:00 2001 From: Anirudh-65 Date: Sat, 7 Nov 2026 16:14:58 +0530 Subject: [PATCH 2/2] [FIX][Relax] Use valid sequence length in test_reverse_sequence --- tests/python/relax/test_frontend_onnx.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 8396db9e1822..c5345bfc6f2f 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -9436,7 +9436,7 @@ def verify_reverse_sequence(x, sequence_lens, time_axis=None, batch_axis=None, o 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, 0, 3], dtype=np.int64)) + 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)