Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent OOB access in QuantizeV2 shape inference
PiperOrigin-RevId: 400309614
Change-Id: I31412c71b05b4f21b677f7fa715a61499cbee39d
  • Loading branch information
miaout17 authored and tensorflower-gardener committed Oct 1, 2021
1 parent 4a7c71d commit a0d6444
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions tensorflow/core/framework/common_shape_fns.cc
Expand Up @@ -2559,6 +2559,9 @@ Status QuantizeV2Shape(InferenceContext* c) {
if (!s.ok() && s.code() != error::NOT_FOUND) {
return s;
}
if (axis < -1) {
return errors::InvalidArgument("axis should be at least -1, got ", axis);
}
const int minmax_rank = (axis == -1) ? 0 : 1;
TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
ShapeHandle minmax;
Expand Down

0 comments on commit a0d6444

Please sign in to comment.