Function make_recv_callback

tensorflow/compiler/jit/kernels/async_io_ops.cc:30–65

namespace {
AsyncIoRendezvous::DoneCallback make_recv_callback(
    OpKernelContext* ctx, AsyncOpKernel::DoneCallback done) {
  using namespace std::placeholders;
  return std::bind(
      [ctx](AsyncOpKernel::DoneCallback done,
            // Begin unbound arguments.
            const Status& s, const AsyncIoRendezvous::TensorPayload& val) {
        Tensor tensor;
        if (val.addr.is_null()) {
          VLOG(2) << "AsyncIoRendezvous::DoneCallback with Tensor size "
                  << val.tensor.TotalBytes();
          tensor = val.tensor;
        } else {
          VLOG(2) << "AsyncIoRendezvous::DoneCallback with payload size "
                  << val.addr.size() << " @" << val.addr.opaque();
          TensorShape tensor_shape;
          XLAShapeToTensorShape(val.shape, &tensor_shape);
          auto data_type =
              EncodePrimitiveTypeAsDataType(val.shape.element_type())
                  .ValueOrDie();
          tensor =
              XlaTensorBuffer::MakeTensor(data_type, tensor_shape, val.addr,
                                          ctx->device()->GetAllocator({}));
        }

        ctx->SetStatus(s);
        if (s.ok()) {
          // 'ctx' allocates the output tensor of the expected type.
          // The runtime checks whether the tensor received here is
          // the same type.
          ctx->set_output(0, tensor);
        }
        done();
      },
      std::move(done), _1, _2);
}
}  // anonymous

XlaAsyncOutSendOp::XlaAsyncOutSendOp(OpKernelConstruction* ctx)