| 80 | } |
| 81 | |
| 82 | Status TensorResponse::ParseFrom(Source* source) { |
| 83 | if (!on_host_) { |
| 84 | protobuf::io::CodedInputStream input(source->contents()); |
| 85 | input.SetTotalBytesLimit(INT_MAX, INT_MAX); // Unlimited |
| 86 | |
| 87 | // Pre-parse into local storage, then delegate to device. |
| 88 | if (!meta_.ParseFromCodedStream(&input) || !input.ConsumedEntireMessage()) { |
| 89 | return errors::InvalidArgument("Cannot parse tensor from response"); |
| 90 | } |
| 91 | Status s = |
| 92 | device_->MakeTensorFromProto(meta_.tensor(), alloc_attrs_, &tensor_); |
| 93 | // Reduce memory usage for big tensors. |
| 94 | { |
| 95 | TensorProto empty; |
| 96 | meta_.mutable_tensor()->Swap(&empty); |
| 97 | } |
| 98 | meta_.clear_tensor(); |
| 99 | return s; |
| 100 | } |
| 101 | if (already_used_) { |
| 102 | ClearTensor(); |
| 103 | } |
| 104 | already_used_ = true; |
| 105 | if (ParseFast(source)) return Status::OK(); |
| 106 | meta_.Clear(); |
| 107 | if (ParseSlow(source)) return Status::OK(); |
| 108 | return errors::InvalidArgument("Cannot parse tensor from response"); |
| 109 | } |
| 110 | |
| 111 | // Define some helper routines for decoding protocol buffer wire format data |
| 112 | namespace { |