`CHECK`-failures during Grappler's `IsSimplifiableReshape` in Tensorflow

@@ -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)) {

@@ -1689,7 +1689,11 @@ Status ConstantFolding::IsSimplifiableReshape(
   if (!IsReshape(node)) {
     return errors::Internal("Node ", node.name(), " is not a Reshape node");
   }
-  CHECK_LE(2, node.input_size());
+  if (2 > node.input_size()) {
+    return errors::Internal("Node ", node.name(),
+                            " must have at most 2 inputs but has ",
+                            node.input_size());
+  }
   const NodeDef* new_shape = node_map_->GetNode(node.input(1));
   if (!IsReallyConstant(*new_shape)) {
     return errors::Internal("Node ", node.name(), " has shape ",
@@ -1707,7 +1711,11 @@ Status ConstantFolding::IsSimplifiableReshape(
   if (!s.ok()) {
     return errors::Internal("Could not evaluate node ", node.name());
   }
-  CHECK_EQ(1, outputs.size());
+  if (outputs.size() != 1) {
+    return errors::Internal("Node ", node.name(),
+                            " must have exactly 1 output but has ",
+                            outputs.size());
+  }
 
   const std::vector<OpInfo::TensorProperties>& props =
       properties.GetInputProperties(node.name());

@@ -1684,15 +1684,17 @@ Status ConstantFolding::FoldGraph(
   return Status::OK();
 }
 
-bool ConstantFolding::IsSimplifiableReshape(
+Status ConstantFolding::IsSimplifiableReshape(
     const NodeDef& node, const GraphProperties& properties) const {
   if (!IsReshape(node)) {
-    return false;
+    return errors::Internal("Node ", node.name(), " is not a Reshape node");
   }
   CHECK_LE(2, node.input_size());
   const NodeDef* new_shape = node_map_->GetNode(node.input(1));
   if (!IsReallyConstant(*new_shape)) {
-    return false;
+    return errors::Internal("Node ", node.name(), " has shape ",
+                            new_shape->DebugString(),
+                            " which is not a constant");
   }
   TensorVector outputs;
   auto outputs_cleanup = gtl::MakeCleanup([&outputs] {
@@ -1703,22 +1705,25 @@ bool ConstantFolding::IsSimplifiableReshape(
 
   Status s = EvaluateNode(*new_shape, TensorVector(), &outputs);
   if (!s.ok()) {
-    return false;
+    return errors::Internal("Could not evaluate node ", node.name());
   }
   CHECK_EQ(1, outputs.size());
 
   const std::vector<OpInfo::TensorProperties>& props =
       properties.GetInputProperties(node.name());
   if (props.empty()) {
-    return false;
+    return errors::Internal("Node ", node.name(), " has no properties");
   }
   const OpInfo::TensorProperties& prop = props[0];
   if (prop.dtype() == DT_INVALID) {
-    return false;
+    return errors::Internal("Node ", node.name(), " has property ",
+                            prop.DebugString(), " with invalid dtype");
   }
   const PartialTensorShape shape(prop.shape());
   if (!shape.IsFullyDefined()) {
-    return false;
+    return errors::Internal("Node ", node.name(), " has property ",
+                            prop.DebugString(), " with shape ",
+                            shape.DebugString(), " which is not fully defined");
   }
 
   PartialTensorShape new_dims;
@@ -1738,7 +1743,12 @@ bool ConstantFolding::IsSimplifiableReshape(
     TF_CHECK_OK(TensorShapeUtils::MakeShape(shp, &new_dims));
   }
 
-  return shape.IsCompatibleWith(new_dims);
+  if (!shape.IsCompatibleWith(new_dims)) {
+    return errors::Internal("Expected shape ", shape.DebugString(),
+                            "to be compatible with ", new_dims.DebugString());
+  }
+
+  return Status::OK();
 }
 
 #define IS_VALUE_CASE(DTYPE, VALUE)                   \
@@ -2925,7 +2935,7 @@ bool ConstantFolding::SimplifyReduction(GraphDef* optimized_graph,
 bool ConstantFolding::SimplifyReshape(const GraphProperties& properties,
                                       bool use_shape_info, NodeDef* node) {
   if (!use_shape_info || node->attr().count("T") == 0 ||
-      !IsSimplifiableReshape(*node, properties)) {
+      !IsSimplifiableReshape(*node, properties).ok()) {
     return false;
   }
   DataType output_type = node->attr().at("T").type();

@@ -129,8 +129,8 @@ class ConstantFolding : public GraphOptimizer {
   Status FoldGraph(const GraphProperties& properties, GraphDef* output,
                    absl::flat_hash_set<string>* nodes_to_not_simplify);
 
-  bool IsSimplifiableReshape(const NodeDef& node,
-                             const GraphProperties& properties) const;
+  Status IsSimplifiableReshape(const NodeDef& node,
+                               const GraphProperties& properties) const;
   Status SimplifyGraph(GraphDef* optimized_graph, GraphProperties* properties,
                        absl::flat_hash_set<string>* nodes_to_not_simplify);
   Status SimplifyNode(NodeDef* node, GraphDef* optimized_graph,