diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index d194b5d55a..11e87edec2 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1142,7 +1142,44 @@ class SymbolicShapeInference: def _infer_SplitToSequence(self, node): self._infer_Split_Common(node, helper.make_sequence_value_info) - def _infer_Squeeze(self, node): + def _infer_Squeeze(self, node): + input_shape = self._get_shape(node, 0) + op_set = get_opset(self.out_mp_) + + # Depending on op-version 'axes' are provided as attribute or via 2nd input + if op_set < 13: + axes = get_attribute(node, 'axes') + assert self._try_get_value(node, 1) is None + else: + axes = self._try_get_value(node, 1) + assert get_attribute(node, 'axes') is None + + if axes is None: + # No axes have been provided (neither via attribute nor via input). + # In this case the 'Shape' op should remove all axis with dimension 1. + # For symbolic dimensions we guess they are !=1. + output_shape = [s for s in input_shape if s != 1] + if self.verbose_ > 0: + symbolic_dimensions = [s for s in input_shape if type(s) != int] + if len(symbolic_dimensions) > 0: + print(f"Symbolic dimensions in input shape of op: '{node.op_type}' node: '{node.name}'. " + + f"Assuming the following dimensions are never equal to 1: {symbolic_dimensions}") + else: + axes = [handle_negative_axis(a, len(input_shape)) for a in axes] + output_shape = [] + for i in range(len(input_shape)): + if i not in axes: + output_shape.append(input_shape[i]) + else: + assert input_shape[i] == 1 or type(input_shape[i]) != int + if self.verbose_ > 0 and type(input_shape[i]) != int: + print(f"Symbolic dimensions in input shape of op: '{node.op_type}' node: '{node.name}'. " + + f"Assuming the dimension '{input_shape[i]}' at index {i} of the input to be equal to 1.") + + vi = self.known_vi_[node.output[0]] + vi.CopyFrom( + helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[0]].type.tensor_type.elem_type, + output_shape)) self._pass_on_sympy_data(node) def _infer_Tile(self, node):