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