Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions onnxoptimizer/passes/pass_util.h
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
25 changes: 25 additions & 0 deletions onnxoptimizer/test/optimizer_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Loading