| 394 | } |
| 395 | |
| 396 | Status TensorHandle::SetRemoteShape(const TensorShape& shape, |
| 397 | tensorflow::Device* d) { |
| 398 | VLOG(3) << "SetRemoteShape on TensorHandle: " << this << " device: " << d; |
| 399 | |
| 400 | if (d != device_) { |
| 401 | mutex_lock l(remote_mirrors_mutex_); |
| 402 | if (remote_mirrors_.find(d) != remote_mirrors_.end()) { |
| 403 | return errors::Internal( |
| 404 | "Attempted to set remote shape for existing mirror."); |
| 405 | } |
| 406 | |
| 407 | auto elem = unshaped_remote_mirrors_.find(d); |
| 408 | if (elem == unshaped_remote_mirrors_.end()) { |
| 409 | return errors::Internal( |
| 410 | "Attempted to set remote shape for non-waiting mirror."); |
| 411 | } |
| 412 | |
| 413 | auto& data = elem->second; |
| 414 | data->ReleaseRemoteTensorHandle(); |
| 415 | remote_mirrors_[d] = absl::make_unique<RemoteTensorHandleData>( |
| 416 | data->op_id(), data->output_num(), shape, data->eager_client(), |
| 417 | data->context_id(), data->ctx()); |
| 418 | unshaped_remote_mirrors_.erase(elem); |
| 419 | |
| 420 | return Status::OK(); |
| 421 | } |
| 422 | |
| 423 | DCHECK(is_remote_) << "SeRemoteShape is only called on remote handles."; |
| 424 | DCHECK(!is_ready_notification_.HasBeenNotified()) |
| 425 | << "SetRemoteShape is only called on non-ready handles."; |
| 426 | |
| 427 | UnshapedRemoteTensorHandleData* p = |
| 428 | reinterpret_cast<UnshapedRemoteTensorHandleData*>( |
| 429 | tensor_handle_data_.get()); |
| 430 | p->ReleaseRemoteTensorHandle(); |
| 431 | tensor_handle_data_ = absl::make_unique<RemoteTensorHandleData>( |
| 432 | remote_op_id_, remote_output_num_, shape, remote_eager_client_, |
| 433 | remote_context_id_, ctx_); |
| 434 | is_poisoned_ = Status::OK(); |
| 435 | is_ready_notification_.Notify(); |
| 436 | |
| 437 | return Status::OK(); |
| 438 | } |
| 439 | #endif |
| 440 | |
| 441 | Status TensorHandle::SetTensor(const tensorflow::Tensor& tensor) { |
no test coverage detected