This is an automated email from the ASF dual-hosted git repository.
tlopex 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 acca1d6dc4 [Fix][Relax][ONNX] Handle Split initializer with
keep_params_in_input (#20091)
acca1d6dc4 is described below
commit acca1d6dc499ea64b552eb003c58dc6866051a9b
Author: Lei Fengxiang <[email protected]>
AuthorDate: Thu Aug 6 03:31:54 2026 +0800
[Fix][Relax][ONNX] Handle Split initializer with keep_params_in_input
(#20091)
Fixes #20066.
When `keep_params_in_input=True`, ONNX initializers are represented as
Relax function parameters while their values remain in the importer
parameter map. `Split._impl_v13` only accepted `relax.Constant`, so it
incorrectly rejected a static initializer as a dynamic Split input.
This change resolves the Split input through the existing `get_constant`
helper before computing the split indices. Genuinely dynamic Split
inputs
remain unsupported and continue to raise the existing error.
The regression test imports an opset-13 model with initializer-backed
split
sizes and verifies both the stored parameter value and the generated
Relax IR.
### Testing
- `cmake --build build --parallel 8`
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_split_initializer_with_params_in_input
tests/python/relax/test_frontend_onnx.py::test_split -xvs` — 2 passed
- `pre-commit run --files
python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 2 +
tests/python/relax/test_frontend_onnx.py | 51 +++++++++++++++++++++++++
2 files changed, 53 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 6d38d2b2ca..d9d97126af 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -2548,6 +2548,8 @@ class Split(OnnxOpConverter):
@classmethod
def _impl_v13(cls, bb, inputs, attr, params):
splits = inputs[1]
+ if splits is not None:
+ splits = get_constant(splits, params)
splits_rank = None
if splits is not None:
splits_rank = splits.ty.ndim
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 730a969b62..a2b7bfb1cd 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -7468,6 +7468,57 @@ def test_split():
)
+def test_split_initializer_with_params_in_input():
+ split_sizes = np.array([2, 4], dtype="int64")
+ split_node = helper.make_node(
+ "Split",
+ ["data", "split_sizes"],
+ ["left", "right"],
+ axis=0,
+ )
+ graph = helper.make_graph(
+ [split_node],
+ "split_initializer_test",
+ inputs=[helper.make_tensor_value_info("data", TensorProto.FLOAT, [6])],
+ initializer=[numpy_helper.from_array(split_sizes, name="split_sizes")],
+ outputs=[
+ helper.make_tensor_value_info("left", TensorProto.FLOAT, [2]),
+ helper.make_tensor_value_info("right", TensorProto.FLOAT, [4]),
+ ],
+ )
+ model = helper.make_model(
+ graph,
+ producer_name="split_initializer_test",
+ opset_imports=[helper.make_opsetid("", 13)],
+ )
+
+ tvm_model = from_onnx(model, opset=13, keep_params_in_input=True)
+ assert len(tvm_model["main"].attrs["params"]) == 1
+
np.testing.assert_array_equal(tvm_model["main"].attrs["params"][0].numpy(),
split_sizes)
+ tvm_model["main"] = tvm_model["main"].without_attr("params")
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ data: R.Tensor((6,), dtype="float32"),
+ split_sizes: R.Tensor((2,), dtype="int64"),
+ ) -> R.Tuple(
+ R.Tensor((2,), dtype="float32"),
+ R.Tensor((4,), dtype="float32"),
+ ):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ lv = R.split(data, indices_or_sections=[2], axis=0)
+ lv1 = lv[0]
+ lv2 = lv[1]
+ gv = (lv1, lv2)
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
def test_tile():
def verify_tile(dynamic, in_shape, repeats, out_shape, expected):
node = helper.make_node("Tile", inputs=["input", "repeats"],
outputs=["out"])
