Fix shape inference for Squeeze without axes #3465
Conversation
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
| @@ -1655,8 +1646,12 @@ ONNX_OPERATOR_SET_SCHEMA( | |||
| return axis < 0 ? axis + input_ndim : axis; | |||
| }); | |||
|
|
|||
| for (int i = 0, j = 0; i < input_ndim; ++i) { | |||
| if (std::find(axes.begin(), axes.end(), i) != axes.end()) { | |||
| for (int i = 0; i < input_ndim; ++i) { | |||
There was a problem hiding this comment.
if-then-else may need to be refactored. If axes is not specified, and input_shape.dim(i) has a symbolic value, it will fall into the else branch, which seems wrong.
There was a problem hiding this comment.
The (not_specify_axes) case is a bit more involved. If there is even one unknown input-dimension, we can't really set any of the other dimensions, since the output rank itself becomes unknown. I suggest creating a temporary shape, which is updated: if we see any unknown-dimension, bail out as we can't even set rank
There was a problem hiding this comment.
Good catch. Since using temporary shape still needs 2-pass for loop (copying from temporary shape needs one), I decided to create another for loop before the original for loop to avoid creating more variables. It will stop early if it encounters a symbolic value from input. Please review it again. Thank you!
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
onnx/defs/tensor/defs.cc
Outdated
| for (int i = 0, j = 0; i < input_ndim; ++i) { | ||
| if (std::find(axes.begin(), axes.end(), i) != axes.end()) { | ||
| for (int i = 0; i < input_ndim; ++i) { | ||
| if (axes_not_specified && input_shape.dim(i).has_dim_value() && |
There was a problem hiding this comment.
nit: input_shape.dim(i).has_dim_value() can be dropped since it is already checked in previous loop.
There was a problem hiding this comment.
Good point. I have removed two of them. Thanks
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
Signed-off-by: Chun-Wei Chen <jacky82226@gmail.com>
Description
Add shape inference for Squeeze without axes: keep dimension if it is not equal to 0.
Motivation
#3456 Currently there is no shape inference for Squeeze without axes.