Method ComputeAsync

tf_euler/kernels/sample_graph_label_op.cc:37–70

};

void SampleGraphLabel::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  auto batch_size = ctx->input(0);
  auto batch_size_flat = batch_size.flat<int32>();
  int32_t batch_size_num = batch_size_flat(0);

  Tensor* output = nullptr;
  TensorShape output_shape;
  output_shape.AddDim(batch_size_num);
  OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &output));

  auto query = new euler::Query(
      "API_SAMPLE_GRAPH_LABEL", "sample_graph", 1,
      {"count"}, {});
  euler::Tensor* count_t = query->AllocInput(
      "count", {1}, euler::DataType::kInt32);
  count_t->Raw<int32_t>()[0] = batch_size_num;

  auto callback = [output, done, query, batch_size_num, this]() {
    euler::Tensor* graph_labels = query->GetResult("sample_graph:0");
    std::string result_s(graph_labels->Raw<char>(),
                         graph_labels->NumElements());
    std::vector<std::string> results_vec = euler::Split(result_s, ",");
    if (results_vec.size() != batch_size_num) {
      EULER_LOG(FATAL) << "results_vec size != batch_size_num";
    }
    auto data = output->flat<tensorflow::string>();
    for (size_t i = 0; i < batch_size_num; ++i) {
      data(i) = results_vec[i];
    }
    delete query;
    done();
  };
  euler::QueryProxy::GetInstance()->RunAsyncGremlin(query, callback);
}

REGISTER_KERNEL_BUILDER(
    Name("SampleGraphLabel").Device(DEVICE_CPU), SampleGraphLabel);