Merge pull request #51330 from tensorflow/mm-cp-d6b57f461b39fd1aa8c1b…

…870f1b974aac3554955-on-r2.4 Prevent nullptr dereference in MLIR TFLite dialect/optimizer.
tensorflow · Aug 6, 2021 · 99e02a9 · 99e02a9
2 parents d843240 + 88075f3
commit 99e02a9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize.cc b/tensorflow/compiler/mlir/lite/transforms/optimize.cc
@@ -61,6 +61,9 @@ constexpr char kRelu6[] = "RELU6";
 constexpr char kRelu1[] = "RELU_N1_TO_1";
 
 bool L2NormalizeReduceAxis(Value sq_op, DenseElementsAttr axis) {
+  if (axis.getNumElements() == 0) {
+    return false;
+  }
   if (sq_op.getType().cast<ShapedType>().getRank() - 1 ==
           *axis.getValues<int>().begin() ||
       *axis.getValues<int>().begin() == -1) {