Skip to content

Commit

Permalink
fix split shape inference error for opset >= 13 (microsoft#19756)
Browse files Browse the repository at this point in the history
### Description
get split operator split section by opset

### Motivation and Context
for opset higher than 13, split section is treated as an input.
  • Loading branch information
inisis authored and Zhenze Wang committed Mar 7, 2024
1 parent 102de3c commit 09c907f
Showing 1 changed file with 11 additions and 2 deletions.
13 changes: 11 additions & 2 deletions onnxruntime/python/tools/symbolic_shape_infer.py
Original file line number Diff line number Diff line change
Expand Up @@ -1940,8 +1940,17 @@ def _infer_SoftmaxCrossEntropyLoss(self, node): # noqa: N802
def _infer_Split_Common(self, node, make_value_info_func): # noqa: N802
input_sympy_shape = self._get_sympy_shape(node, 0)
axis = handle_negative_axis(get_attribute(node, "axis", 0), len(input_sympy_shape))
split = get_attribute(node, "split")
if not split:
op_set = get_opset(self.out_mp_)

# Depending on op-version 'split' are provided as attribute or via 2nd input
if op_set < 13:
split = get_attribute(node, "split")
assert self._try_get_value(node, 1) is None
else:
split = self._try_get_value(node, 1)
assert get_attribute(node, "split") is None

if split is None:
num_outputs = len(node.output)
split = [input_sympy_shape[axis] / sympy.Integer(num_outputs)] * num_outputs
self._update_computed_dims(split)
Expand Down

0 comments on commit 09c907f

Please sign in to comment.