handle_status = computation->AddInfeedInstruction(arg->infeed_request());
break;
case OpRequest::kOutfeedRequest:
- TF_RETURN_IF_ERROR(
- computation->AddOutfeedInstruction(arg->outfeed_request()));
- return tensorflow::Status::OK();
+ handle_status =
+ computation->AddOutfeedInstruction(arg->outfeed_request());
+ break;
case OpRequest::kMapRequest: {
TF_ASSIGN_OR_RETURN(
UserComputation * to_apply,
return handle;
}
-Status UserComputation::AddOutfeedInstruction(
+StatusOr<ComputationDataHandle> UserComputation::AddOutfeedInstruction(
const OutfeedRequest& outfeed_request) {
tensorflow::mutex_lock lock(mutex_);
// Verify that operand is valid.
TF_RETURN_IF_ERROR(LookUpRequest(outfeed_request.operand()).status());
- // No handle is returned, but a handle must be assigned to this instruction
- // for computation versioning.
ComputationDataHandle handle = CreateComputationDataHandle();
OperationRequest& request =
(*session_computation_.mutable_requests())[handle.handle()];
VLOG(1) << "AddOutfeedInstruction (" << GetVersionedHandleInternal()
<< "), data handle " << handle.handle() << ": "
<< outfeed_request.ShortDebugString();
- return Status::OK();
+ return handle;
}
StatusOr<ComputationDataHandle> UserComputation::AddCallInstruction(
const InfeedRequest& infeed_request);
// Enqueues an outfeed instruction onto this user computation.
- Status AddOutfeedInstruction(const OutfeedRequest& outfeed_request);
+ StatusOr<ComputationDataHandle> AddOutfeedInstruction(
+ const OutfeedRequest& outfeed_request);
// Enqueues a call instruction onto this user computation.
StatusOr<ComputationDataHandle> AddCallInstruction(
*outfeed_request.mutable_operand() = constant_handle;
*outfeed_request.mutable_shape() = kVectorShape;
outfeed_request.set_outfeed_config("abc");
- TF_ASSERT_OK(computation.AddOutfeedInstruction(outfeed_request));
+ TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle outfeed_handle,
+ computation.AddOutfeedInstruction(outfeed_request));
auto hlo_resolver = [](const VersionedComputationHandle& handle) {
return nullptr;