saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix shape inference bug (#19848)

Browse source 
### Description
for nodes like add, their input should be merged dynamically

### Motivation and Context
when doing shape inference, for nodes like Add, currently when doing _onnx_infer_single_node, their inputs are generated from last node's output, but they should be merged.
This commit is contained in:
inisis 2024-03-30 04:06:27 +08:00 committed by GitHub
parent b1a5eb255e
commit 8396845806
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 22 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -495,6 +495,28 @@ class SymbolicShapeInference:
if (name in self.initializers_ and name not in self.graph_inputs_)
]
if node.op_type in [
"Add",
"Sub",
"Mul",
"Div",
"MatMul",
"MatMulInteger",
"MatMulInteger16",
"Where",
"Sum",
]:
if node.output[0] in self.known_vi_:
vi = self.known_vi_[node.output[0]]
out_rank = len(get_shape_from_type_proto(vi.type))
in_shapes = [self._get_shape(node, i) for i in range(len(node.input))]
for d in range(
out_rank - (2 if node.op_type in ["MatMul", "MatMulInteger", "MatMulInteger16"] else 0)
):
in_dims = [s[len(s) - out_rank + d] for s in in_shapes if len(s) + d >= out_rank]
if len(in_dims) > 1:
self._check_merged_dims(in_dims, allow_broadcast=True)
# run single node inference with self.known_vi_ shapes
tmp_graph = helper.make_graph(
[node],