| 370 | } |
| 371 | |
| 372 | Status SliceRecvOp::RecvShape(OpKernelContext* ctx, |
| 373 | const FrameAndIter& frame_iter, |
| 374 | TensorShape& shape) { |
| 375 | Rendezvous::Args args; |
| 376 | args.device_context = ctx->op_device_context(); |
| 377 | args.alloc_attrs = AllocatorAttributes(); |
| 378 | if (ctx->is_eager()) { |
| 379 | // NOTE(fishx): Only set cancellation_manager in eager mode. Because in |
| 380 | // Tensorflow 1.x, session (or graph_mgr) will abort the underlying |
| 381 | // rendezvous if it encounters any error. |
| 382 | args.cancellation_manager = ctx->cancellation_manager(); |
| 383 | } |
| 384 | |
| 385 | Rendezvous::ParsedKey parsed_key; |
| 386 | slice_sendrecv::GetSliceRendezvousKey(key_prefix_, "_slice_transfer_shape", |
| 387 | frame_iter, &parsed_key.buf_); |
| 388 | VLOG(2) << "SliceRecv " << parsed_key.buf_; |
| 389 | TF_RETURN_IF_ERROR(Rendezvous::ParseKey(parsed_key.buf_, &parsed_key)); |
| 390 | |
| 391 | Tensor shape_t; |
| 392 | bool is_dead; |
| 393 | TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv(parsed_key, args, &shape_t, |
| 394 | &is_dead, timeout_ms_)); |
| 395 | // This shouldn't be a dead tensor. |
| 396 | CHECK_EQ(is_dead, false); |
| 397 | auto shape_vec = shape_t.vec<int64>(); |
| 398 | const int64 num_elements = shape_t.NumElements(); |
| 399 | for (int64 i = 0; i < num_elements; i++) { |
| 400 | shape.AddDim(shape_vec(i)); |
| 401 | } |
| 402 | |
| 403 | return Status::OK(); |
| 404 | } |
| 405 | |
| 406 | Status SliceRecvOp::RecvString(OpKernelContext* ctx, |
| 407 | const FrameAndIter& frame_iter, |
nothing calls this directly
no test coverage detected