Minor code cleanup in remote tensors for better readability.

PiperOrigin-RevId: 291215374
Change-Id: Ib6336634eb464c56bcaf636d36226649c2a40a97
This commit is contained in:
Haoyu Zhang 2020-01-23 12:01:44 -08:00 committed by TensorFlower Gardener
parent 6a233b048b
commit e42a740859
2 changed files with 10 additions and 5 deletions

View File

@ -95,7 +95,8 @@ Status TensorHandle::GetResourceHandleDtypesAndShapes(
Status TensorHandle::CreateLocalHandle(const class Tensor& t,
TensorHandle** h) {
// TODO(b/136608821): Move away from nullptr
return CreateLocalHandle(t, nullptr, nullptr, nullptr, h);
return CreateLocalHandle(t, /*d=*/nullptr, /*op_device=*/nullptr,
/*ctx=*/nullptr, h);
}
Status TensorHandle::CreateLocalHandle(const class Tensor& t, Device* d,

View File

@ -56,7 +56,7 @@ Status CreateUncachedKernelAndDeviceOp(
ctx.HostCPU()));
const NodeDef& ndef = op->MutableAttrs()->BuildNodeDef();
return kernel->get()->Init(ndef, nullptr);
return kernel->get()->Init(ndef, /*graph_collector=*/nullptr);
}
// This gets a unique wire ID. We add a random identifier so that if the
@ -105,14 +105,17 @@ Status RemoteCopyNode::RunLocalSend(EagerOperation* op) {
TF_RETURN_IF_ERROR(src_->TensorValue(&input_vector[0]));
EagerKernelArgs args(std::move(input_vector));
return kernel->Run(args, nullptr, nullptr, absl::nullopt);
return kernel->Run(args, /*outputs=*/nullptr,
/*cancellation_manager=*/nullptr,
/*remote_func_params=*/absl::nullopt);
}
void RemoteCopyNode::StartSend() {
// TODO(gjn): We should consider just using the low-level SendOp::Compute()
// functionality here instead of constructing an Op.
EagerOperation op(ctx_);
Status status = op.Reset("_Send", nullptr, false, nullptr);
Status status = op.Reset("_Send", /*raw_device_name=*/nullptr,
/*remote=*/false, /*executor=*/nullptr);
if (!status.ok()) {
captured_state_->SetSendStatus(status);
return;
@ -244,7 +247,8 @@ void RemoteCopyNode::StartRecv(StatusCallback done) {
// TODO(gjn): We should consider just using the low-level RecvOp::Compute()
// functionality here instead of constructing an Op.
EagerOperation op(ctx_);
Status status = op.Reset("_Recv", nullptr, false, nullptr);
Status status = op.Reset("_Recv", /*raw_device_name=*/nullptr,
/*remote=*/false, /*executor=*/nullptr);
if (!status.ok()) {
captured_state_->dst()->Poison(status);
done(status);