tensor_bundle: tighten Lookup() semantics, add LookupDtypeShape().

This fixes a subtle allocation issue in CheckpointReader.
Change: 135503615

1 files changed, 8 insertions, 3 deletions

diff --git a/tensorflow/c/checkpoint_reader.cc b/tensorflow/c/checkpoint_reader.cc
index 23a42cffbb..2ac3f75e5b 100644
--- a/tensorflow/c/checkpoint_reader.cc
+++ b/tensorflow/c/checkpoint_reader.cc
@@ -81,9 +81,14 @@ void CheckpointReader::GetTensor(
   if (reader_ != nullptr) {
     status = reader_->GetTensor(name, out_tensor);
   } else {
-    std::unique_ptr<Tensor> tensor(new Tensor);
-    status = v2_reader_->Lookup(name, tensor.get());
-    if (status.ok()) std::swap(*out_tensor, tensor);
+    tensorflow::DataType dtype;
+    tensorflow::TensorShape shape;
+    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
+    if (status.ok()) {
+      out_tensor->reset(new Tensor(dtype, shape));
+      status = v2_reader_->Lookup(name, out_tensor->get());
+      if (!status.ok()) out_tensor->reset();
+    }
   }
   if (!status.ok()) {
     Set_TF_Status_from_Status(out_status, status);