From 04b3a9f1246ac15dbbfc05501a20d46538c362cf Mon Sep 17 00:00:00 2001 From: Yuhx141 <155034285+Yuhx141@users.noreply.github.com> Date: Thu, 1 Oct 2026 14:33:17 +0800 Subject: [PATCH] Clamp Shape slice bounds during elimination Signed-off-by: Yuhx141 <155034285+Yuhx141@users.noreply.github.com> --- onnxoptimizer/passes/pass_util.h | 14 +++++++++----- onnxoptimizer/test/optimizer_test.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/onnxoptimizer/passes/pass_util.h b/onnxoptimizer/passes/pass_util.h index 759127084..a805033d4 100644 --- a/onnxoptimizer/passes/pass_util.h +++ b/onnxoptimizer/passes/pass_util.h @@ -320,11 +320,15 @@ inline std::pair FetchStartAndEndAttrOfShape( const Node* shape, const int64_t rank) { ONNX_ASSERT(CheckKind(shape, "Shape")); - const int64_t start = AddYIfNegative( - GetValueFromAttrWithDefault(shape, "start", 0), rank); - const int64_t end = AddYIfNegative( - GetValueFromAttrWithDefault(shape, "end", rank), rank); - return {start, end}; + const int64_t start = std::clamp( + AddYIfNegative(GetValueFromAttrWithDefault(shape, "start", 0), + rank), + 0, rank); + const int64_t end = std::clamp( + AddYIfNegative(GetValueFromAttrWithDefault(shape, "end", rank), + rank), + 0, rank); + return {start, std::max(start, end)}; } inline std::pair FetchStartAndEndAttrOfShape( diff --git a/onnxoptimizer/test/optimizer_test.py b/onnxoptimizer/test/optimizer_test.py index 1800f1242..b96ac7898 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/onnxoptimizer/test/optimizer_test.py @@ -3652,6 +3652,34 @@ def test_eliminate_shape_op(self): # type: () -> None assert len(optimized_model.graph.node) == 1 assert optimized_model.graph.node[0].op_type == "Relu" + def test_eliminate_shape_op_clamps_slice_bounds(self): # type: () -> None + for start, end, expected in [(-5, 10, [3, 2]), (5, 7, []), (1, 0, [])]: + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 2]) + X2 = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [3, 2]) + Y = helper.make_tensor_value_info("Y", TensorProto.INT64, [len(expected)]) + graph = helper.make_graph( + [ + helper.make_node("Relu", ["X"], ["X2"]), + helper.make_node("Shape", ["X2"], ["Y"], start=start, end=end), + ], + "test", + [X], + [Y], + value_info=[X2], + ) + + optimized_model = self._optimized( + graph, + ["eliminate_shape_op"], + False, + compare_result=False, + opset_imports=[helper.make_opsetid("", 15)], + ) + + assert [node.op_type for node in optimized_model.graph.node] == ["Relu"] + assert len(optimized_model.graph.initializer) == 1 + assert to_array(optimized_model.graph.initializer[0]).tolist() == expected + # Some exporters use negative dim to represent dynamic shapes. def test_not_eliminate_shape_op_with_negative_dim(self): # type: () -> None X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, -1])