| 768 | } |
| 769 | |
| 770 | Status run(BundleReader* reader) { |
| 771 | TensorShape restored_full_shape; |
| 772 | TF_RETURN_IF_ERROR( |
| 773 | reader->LookupTensorShape(tensor_name, &restored_full_shape)); |
| 774 | |
| 775 | VLOG(1) << "Restoring tensor " << idx << " : " << tensor_name << " : " |
| 776 | << restored_full_shape.num_elements(); |
| 777 | Tensor* restored_tensor; |
| 778 | if (shape_and_slice.empty()) { |
| 779 | // Lookup the full tensor. |
| 780 | TF_RETURN_IF_ERROR( |
| 781 | context->allocate_output(idx, restored_full_shape, &restored_tensor)); |
| 782 | TF_RETURN_IF_ERROR(reader->Lookup(tensor_name, restored_tensor)); |
| 783 | } else { |
| 784 | // Lookup the slice. |
| 785 | TensorShape parsed_full_shape; |
| 786 | TensorSlice parsed_slice; |
| 787 | TensorShape parsed_slice_shape; |
| 788 | |
| 789 | TF_RETURN_IF_ERROR( |
| 790 | checkpoint::ParseShapeAndSlice(shape_and_slice, &parsed_full_shape, |
| 791 | &parsed_slice, &parsed_slice_shape)); |
| 792 | |
| 793 | if (!restored_full_shape.IsSameSize(parsed_full_shape)) { |
| 794 | return errors::InvalidArgument( |
| 795 | "tensor_name = ", tensor_name, "; shape in shape_and_slice spec ", |
| 796 | parsed_full_shape.DebugString(), |
| 797 | " does not match the shape stored in checkpoint: ", |
| 798 | restored_full_shape.DebugString()); |
| 799 | } |
| 800 | TF_RETURN_IF_ERROR( |
| 801 | context->allocate_output(idx, parsed_slice_shape, &restored_tensor)); |
| 802 | TF_RETURN_IF_ERROR( |
| 803 | reader->LookupSlice(tensor_name, parsed_slice, restored_tensor)); |
| 804 | } |
| 805 | return Status::OK(); |
| 806 | } |
| 807 | |
| 808 | OpKernelContext* context; |
| 809 | size_t idx; |
no test coverage detected