diff --git a/onnxoptimizer/passes/pass_util.h b/onnxoptimizer/passes/pass_util.h index 759127084..49701f08b 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..4ad13bbae 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/onnxoptimizer/test/optimizer_test.py @@ -3652,6 +3652,31 @@ 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_start_beyond_rank(self): # type: () -> None + X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 2]) + X2 = helper.make_tensor_value_info("X2", TensorProto.INT64, [0]) + Y = helper.make_tensor_value_info("Y", TensorProto.INT64, [0]) + graph = helper.make_graph( + [ + helper.make_node("Shape", ["X"], ["X2"], start=5), + helper.make_node("Identity", ["X2"], ["Y"]), + ], + "test", + [X], + [Y], + value_info=[X2], + ) + optimized_model = self._optimized( + graph, + ["eliminate_shape_op"], + False, + opset_imports=[helper.make_opsetid("", 15)], + ) + + assert len(optimized_model.graph.node) == 1 + assert optimized_model.graph.node[0].op_type == "Identity" + assert numpy_helper.to_array(optimized_model.graph.initializer[0]).tolist() == [] + # 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])