| 551 | } |
| 552 | |
| 553 | TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { |
| 554 | if (h == nullptr || h->handle == nullptr) { |
| 555 | status->status = tensorflow::errors::InvalidArgument( |
| 556 | "The passed in handle is a nullptr"); |
| 557 | return nullptr; |
| 558 | } |
| 559 | tensorflow::TensorHandle* handle = h->handle; |
| 560 | |
| 561 | // TODO(agarwal): move this implementation inside TFE_TensorHandle. |
| 562 | if (handle->IsRemote()) { |
| 563 | const tensorflow::Tensor* t = nullptr; |
| 564 | tensorflow::TensorHandle* h_cpu = nullptr; |
| 565 | status->status = EagerCopyToDevice( |
| 566 | handle, handle->Context(), handle->Context()->Executor(), |
| 567 | handle->Context()->HostCPU(), false, &h_cpu); |
| 568 | if (!status->status.ok()) { |
| 569 | return nullptr; |
| 570 | } |
| 571 | status->status = h_cpu->Tensor(&t); |
| 572 | if (!status->status.ok()) { |
| 573 | h_cpu->Unref(); |
| 574 | return nullptr; |
| 575 | } |
| 576 | TF_Tensor* retval = tensorflow::TF_TensorFromTensor(*t, status); |
| 577 | h_cpu->Unref(); |
| 578 | return retval; |
| 579 | } else { |
| 580 | tensorflow::Tensor tensor; |
| 581 | if (IsCPU(handle->device())) { |
| 582 | const tensorflow::Tensor* src = nullptr; |
| 583 | status->status = handle->Tensor(&src); |
| 584 | if (!status->status.ok()) return nullptr; |
| 585 | tensor = *src; |
| 586 | } else { |
| 587 | tensorflow::EagerContext* ctx = handle->Context(); |
| 588 | CHECK_NE(ctx, nullptr); |
| 589 | status->status = h->handle->CopyToDevice(ctx, ctx->HostCPU(), &tensor); |
| 590 | if (!status->status.ok()) return nullptr; |
| 591 | } |
| 592 | return tensorflow::TF_TensorFromTensor(tensor, status); |
| 593 | } |
| 594 | } |
| 595 | |
| 596 | TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, |
| 597 | TF_Status* status) { |