From 301023d9de421eb57db9466ae4044273f1783148 Mon Sep 17 00:00:00 2001 From: Javier de Jesus Date: Sat, 13 Jun 2026 19:37:38 +0000 Subject: [PATCH] [Relax][ONNX] Support exclusive option in CumSum The CumSum converter aborted on any node with exclusive=1 via a bare assert. relax.op.cumsum and topi.cumsum already implement the exclusive form, so pass the attribute through instead of rejecting it. The existing flip-cumsum-flip reverse handling composes correctly with exclusive. Extend test_cumsum to cover exclusive=True so check_correctness exercises all four reverse/exclusive combinations against ONNX Runtime. Fixes #19692 --- python/tvm/relax/frontend/onnx/onnx_frontend.py | 4 ++-- tests/python/relax/test_frontend_onnx.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index 3d9dfba9a16b..5be4f0831d90 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py @@ -1905,7 +1905,7 @@ class CumSum(OnnxOpConverter): def _impl_v14(cls, bb, inputs, attr, params): data = inputs[0] axis_input = get_constant(inputs[1], params) - assert not attr.get("exclusive", False), "Exclusive option not yet supported." + exclusive = attr.get("exclusive", 0) != 0 if isinstance(axis_input, relax.Constant): axis_data = axis_input.data.numpy() @@ -1933,7 +1933,7 @@ def _impl_v14(cls, bb, inputs, attr, params): if attr.get("reverse", 0) != 0: data = bb.emit_te(topi.flip, data, axis=axis) - data = relax.op.cumsum(data, axis) + data = relax.op.cumsum(data, axis, exclusive=exclusive) data = bb.normalize(data) if attr.get("reverse", 0) != 0: diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 8d8c1bc54b9b..0db3739daa85 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -2001,7 +2001,7 @@ def test_pow(): @pytest.mark.parametrize("reverse", [True, False]) -@pytest.mark.parametrize("exclusive", [False]) +@pytest.mark.parametrize("exclusive", [True, False]) def test_cumsum(reverse, exclusive): cumsum_node = helper.make_node( "CumSum", ["x", "axis"], ["y"], reverse=reverse, exclusive=exclusive