Skip to content

[Bug][Relax][ONNX] CumSum with exclusive=1 hits a bare AssertionError #19692

@wuyii8941

Description

@wuyii8941

Expected behavior

ONNX CumSum has two boolean attributes — reverse and exclusive. With exclusive=1, the output excludes the current element:

Y[..., i] = sum(X[..., 0..i-1]) (exclusive=1, reverse=0)

ONNX Runtime supports both attributes; this is the standard PyTorch cumsum semantics shifted by one.

Actual behavior

The TVM frontend hits a bare assertion that fails the entire from_onnx call:

AssertionError: Exclusive option not yet supported.

This is not even raised as a structured NotImplementedError — it's an assert with no operator/opset context, making downstream error reporting confusing.

Reproduction

import numpy as np
import onnx
from onnx import helper, TensorProto, numpy_helper
import onnxruntime as ort
from tvm.relax.frontend.onnx import from_onnx

X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4])
Y = onnx.ValueInfoProto(); Y.name = "Y"
axis = numpy_helper.from_array(np.array(1, dtype=np.int64), "axis")
node = helper.make_node("CumSum", ["X", "axis"], ["Y"], exclusive=1, reverse=0)
g = helper.make_graph([node], "g", [X], [Y], initializer=[axis])
m = helper.make_model(g, opset_imports=[helper.make_opsetid("", 14)])

arr = np.array([[1., 2., 3., 4.]], dtype=np.float32)
print("ORT:", ort.InferenceSession(m.SerializeToString()).run(None, {"X": arr})[0])
# ORT: [[0. 1. 3. 6.]]

inf = onnx.shape_inference.infer_shapes(m)
mod = from_onnx(inf)  # AssertionError: Exclusive option not yet supported.

Root cause

python/tvm/relax/frontend/onnx/onnx_frontend.py, CumSum._impl_v14:

@classmethod
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."
    ...

topi.cumsum (and relax.op.cumsum) only implement the inclusive form, but the converter just asserts instead of (a) raising a structured error or (b) emitting an exclusive decomposition.

Suggested fix

Two-stage approach:

  1. Short-term: replace the assertion with a clear NotImplementedError that mentions the op and opset:

    if attr.get("exclusive", False):
        raise NotImplementedError(
            "CumSum with exclusive=1 is not yet supported in the Relax ONNX frontend"
        )
  2. Better: decompose exclusive to inclusive cumsum followed by a one-element shift, since this is straightforward:

    y = relax.op.cumsum(data, axis=axis)
    if exclusive:
        # shift right by one along `axis`, fill with identity (0)
        y = relax.op.subtract(y, data)

reverse=1 already lacks a dedicated branch and silently falls through to relax.op.cumsum(...) — that should be audited as well.

Impact

Any model using ONNX CumSum with exclusive=1 cannot be imported. Common in attention masking utilities and statistical models.

Environment

  • TVM: latest main (commit b172d5e)
  • Python: 3.11
  • ONNX Runtime: 1.24.4

cc @KJlaccHoeUM9l @junrushao

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions