Expr propagation bug.

tensorflow/core/framework/tensor_shape.cc

@@ -496,8 +496,11 @@ void TensorShapeRep::set_expression(int d, xla::DynExpr* expr) {
     expressions_.clear();
     return;
   }
-  if (expressions_.size() <= static_cast<size_t>(d)) {
-    expressions_.resize(d + 1, nullptr);
+  CHECK_GE(d, 0);
+  CHECK_LT(d, ndims_byte());
+  const size_t new_size = static_cast<size_t>(d) + 1;
+  if (expressions_.size() < new_size) {
+    expressions_.resize(new_size, nullptr);
   }
   expressions_[d] = expr;
 }
@@ -515,6 +518,9 @@ void TensorShapeRep::set_expressions(std::vector<xla::DynExpr*> exprs) {
     expressions_.clear();
     return;
   }
+  if (exprs.size() > ndims_byte()) {
+    exprs.resize(ndims_byte());
+  }
   while (!exprs.empty() && exprs.back() == nullptr) {
     exprs.pop_back();
   }

tensorflow/core/framework/tensor_shape.h

@@ -84,7 +84,9 @@ class TensorShapeRep {
 
   // Get the array of dynamic multipliers.
   std::vector<xla::DynExpr*> get_expressions() const {
-    return expressions_;
+    const size_t limit = std::min<size_t>(expressions_.size(), ndims_byte());
+    return std::vector<xla::DynExpr*>(expressions_.begin(),
+                                      expressions_.begin() + limit);
   }
 
   // Get the array of dynamic multipliers, filling missing entries with

tensorflow/core/grappler/costs/graph_properties.cc

@@ -2105,7 +2105,9 @@ class SymbolicShapeRefiner {
       }
     }
     // Update NodeContext output fields after shape inference function runs.
-    status.Update(CanonicalizeOutputDims(&node));
+    if (TensorShapeExpressionsEnabled()) {
+      status.Update(CanonicalizeOutputDims(&node));
+    }
     status.Update(MaybeUpdateNodeContextOutput(node, is_fed, c));
 
     return status;