| 958 | } |
| 959 | |
| 960 | Status Service::TransferToClient(const TransferToClientRequest* arg, |
| 961 | TransferToClientResponse* result) { |
| 962 | TF_ASSIGN_OR_RETURN(const ShapedBuffer* shaped_buffer, |
| 963 | allocation_tracker_.ResolveForReplica(arg->data(), 0)); |
| 964 | |
| 965 | Shape return_shape; |
| 966 | if (arg->has_shape_with_layout()) { |
| 967 | return_shape = Shape(arg->shape_with_layout()); |
| 968 | if (!LayoutUtil::HasLayout(return_shape)) { |
| 969 | return InvalidArgument("shape_with_layout must have layout if present."); |
| 970 | } |
| 971 | } else { |
| 972 | return_shape = Shape(shaped_buffer->on_host_shape()); |
| 973 | } |
| 974 | |
| 975 | TF_ASSIGN_OR_RETURN(auto stream, execute_backend_->BorrowStream( |
| 976 | shaped_buffer->device_ordinal())); |
| 977 | |
| 978 | TF_ASSIGN_OR_RETURN( |
| 979 | Literal result_literal, |
| 980 | execute_backend_->transfer_manager()->TransferLiteralFromDevice( |
| 981 | stream.get(), *shaped_buffer)); |
| 982 | |
| 983 | if (LayoutUtil::LayoutsInShapesEqual(return_shape, result_literal.shape())) { |
| 984 | *result->mutable_literal() = result_literal.ToProto(); |
| 985 | } else { |
| 986 | *result->mutable_literal() = |
| 987 | result_literal.Relayout(return_shape).ToProto(); |
| 988 | } |
| 989 | return Status::OK(); |
| 990 | } |
| 991 | |
| 992 | Status Service::TransferToServer(const TransferToServerRequest* arg, |
| 993 | TransferToServerResponse* result) { |
nothing calls this directly
no test coverage detected