Clamp Shape slice bounds before eliminating Shape

onnxoptimizer/passes/pass_util.h

@@ -320,11 +320,15 @@ inline std::pair<int64_t, int64_t> FetchStartAndEndAttrOfShape(
     const Node* shape, const int64_t rank) {
   ONNX_ASSERT(CheckKind(shape, "Shape"));
 
-  const int64_t start = AddYIfNegative<int64_t>(
-      GetValueFromAttrWithDefault(shape, "start", 0), rank);
-  const int64_t end = AddYIfNegative<int64_t>(
-      GetValueFromAttrWithDefault(shape, "end", rank), rank);
-  return {start, end};
+  const int64_t start = std::clamp<int64_t>(
+      AddYIfNegative<int64_t>(
+          GetValueFromAttrWithDefault(shape, "start", 0), rank),
+      0, rank);
+  const int64_t end = std::clamp<int64_t>(
+      AddYIfNegative<int64_t>(
+          GetValueFromAttrWithDefault(shape, "end", rank), rank),
+      0, rank);
+  return {start, std::max(start, end)};
 }
 
 inline std::pair<int64_t, int64_t> FetchStartAndEndAttrOfShape(

onnxoptimizer/test/ir_test.py

@@ -3,14 +3,32 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import onnx
-import onnx.backend.test.loader
 
 import onnxoptimizer
 
 
 def test_seq_type():
-    orig = onnx.load(
-        onnx.backend.test.loader.DATA_DIR + "/node/test_sequence_insert_at_back/model.onnx"
+    sequence = onnx.helper.make_tensor_sequence_value_info(
+        "sequence", onnx.TensorProto.INT64, None
+    )
+    tensor = onnx.helper.make_tensor_value_info(
+        "tensor", onnx.TensorProto.INT64, [3]
+    )
+    output = onnx.helper.make_tensor_sequence_value_info(
+        "output_sequence", onnx.TensorProto.INT64, None
+    )
+    orig = onnx.helper.make_model(
+        onnx.helper.make_graph(
+            [
+                onnx.helper.make_node(
+                    "SequenceInsert", ["sequence", "tensor"], ["output_sequence"]
+                )
+            ],
+            "sequence_insert",
+            [sequence, tensor],
+            [output],
+        ),
+        opset_imports=[onnx.helper.make_opsetid("", 11)],
     )
     optimized = onnxoptimizer.optimize(orig)
     assert optimized.graph.input[0].type == orig.graph.input[0].type

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])