Eliminate CHECK-fails from IsSimplifiableReshape via MakeShape(<invalid shape>)

mihaimaruseac · tensorflower-gardener · commit 240655511cd3 · 2021-11-11T09:29:48.000-08:00
PiperOrigin-RevId: 409166738
Change-Id: I7f0a3590b8acae3f3e3e2fe636e1f5ef285693cf
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.cc b/tensorflow/core/grappler/optimizers/constant_folding.cc
@@ -1741,14 +1741,16 @@ Status ConstantFolding::IsSimplifiableReshape(
       int32_t dim = outputs[0]->flat<int32>()(i);
       shp.push_back(dim);
     }
-    TF_CHECK_OK(TensorShapeUtils::MakeShape(shp, &new_dims));
+    s = TensorShapeUtils::MakeShape(shp, &new_dims);
+    if (!s.ok()) return s;
   } else {
     std::vector<int64_t> shp;
     for (int i = 0; i < outputs[0]->NumElements(); ++i) {
       int64_t dim = outputs[0]->flat<int64_t>()(i);
       shp.push_back(dim);
     }
-    TF_CHECK_OK(TensorShapeUtils::MakeShape(shp, &new_dims));
+    s = TensorShapeUtils::MakeShape(shp, &new_dims);
+    if (!s.ok()) return s;
   }
 
   if (!shape.IsCompatibleWith(new_dims)) {

Original file line number	Diff line number	Diff line change
`@@ -1741,14 +1741,16 @@ Status ConstantFolding::IsSimplifiableReshape(`
`1741`	`1741`	`int32_t dim = outputs[0]->flat<int32>()(i);`
`1742`	`1742`	`shp.push_back(dim);`
`1743`	`1743`	`}`
`1744`		`- TF_CHECK_OK(TensorShapeUtils::MakeShape(shp, &new_dims));`
	`1744`	`+ s = TensorShapeUtils::MakeShape(shp, &new_dims);`
	`1745`	`+ if (!s.ok()) return s;`
`1745`	`1746`	`} else {`
`1746`	`1747`	`std::vector<int64_t> shp;`
`1747`	`1748`	`for (int i = 0; i < outputs[0]->NumElements(); ++i) {`
`1748`	`1749`	`int64_t dim = outputs[0]->flat<int64_t>()(i);`
`1749`	`1750`	`shp.push_back(dim);`
`1750`	`1751`	`}`
`1751`		`- TF_CHECK_OK(TensorShapeUtils::MakeShape(shp, &new_dims));`
	`1752`	`+ s = TensorShapeUtils::MakeShape(shp, &new_dims);`
	`1753`	`+ if (!s.ok()) return s;`
`1752`	`1754`	`}`
`1753`	`1755`
`1754`	`1756`	`if (!shape.IsCompatibleWith(new_dims)) {`