Merge pull request #51292 from tensorflow/mm-cherrypick-da857cfa0fde8…

…f79ad0afdbc94e88b5d4bbec764-on-r2.3 Fix a shape inference issue leading to nullptr deref.
tensorflow · Aug 5, 2021 · 6ca0352 · 6ca0352
2 parents b43738a + ad9bb3e
commit 6ca0352
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
@@ -2887,6 +2887,10 @@ REGISTER_OP("Dequantize")
       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;