Method ProcessBatch

tensorflow/core/kernels/batch_kernels.cc:567–635

Processes a batch of one or more BatchTask entries.

  // Processes a batch of one or more BatchTask entries.
  void ProcessBatch(std::unique_ptr<Batch> batch) const {
    if (batch->empty()) {
      return;
    }

    WithContext wc(batch->task(batch->num_tasks() - 1).propagated_context);

    OpKernelContext* last_task_context =
        batch->task(batch->num_tasks() - 1).context;
    AsyncOpKernel::DoneCallback last_task_callback =
        batch->task(batch->num_tasks() - 1).done_callback;

    OP_REQUIRES_OK_ASYNC(last_task_context, ValidateBatch(*batch),
                         last_task_callback);

    // All tasks should have the same number of input edges.
    const int num_input_edges = batch->task(0).inputs.size();
    std::vector<Tensor> concatenated_tensors;
    const Status concat_status =
        ConcatInputTensors(*batch, last_task_context, &concatenated_tensors);
    OP_REQUIRES_OK_ASYNC(last_task_context, concat_status, last_task_callback);

    // Process each input edge one at a time (the typical case has just one).
    for (int i = 0; i < num_input_edges; ++i) {
      last_task_context->set_output(i, concatenated_tensors.at(i));

      // Emit batch->num_tasks() - 1 empty output tensors.
      for (int task_idx = 0; task_idx < batch->num_tasks() - 1; ++task_idx) {
        const BatchTask& task = batch->task(task_idx);
        TensorShape output_shape(task.inputs.at(i).shape());
        output_shape.set_dim(0, 0);
        Tensor* output = nullptr;
        OP_REQUIRES_OK_ASYNC(
            task.context,
            task.context->allocate_output(i, output_shape, &output),
            task.done_callback);
      }
    }
    // Emit batch->num_tasks() - 1 empty index tensors.
    for (int task_idx = 0; task_idx < batch->num_tasks() - 1; ++task_idx) {
      const BatchTask& task = batch->task(task_idx);
      TensorShape index_shape({0, 3});
      Tensor* output = nullptr;
      OP_REQUIRES_OK_ASYNC(
          task.context,
          task.context->allocate_output(num_input_edges, index_shape, &output),
          task.done_callback);
    }
    // Emit all ID tensors.
    for (int task_idx = 0; task_idx < batch->num_tasks(); ++task_idx) {
      const BatchTask& task = batch->task(task_idx);
      Tensor* id;
      OP_REQUIRES_OK_ASYNC(task.context,
                           task.context->allocate_output(num_input_edges + 1,
                                                         TensorShape({}), &id),
                           task.done_callback);
      id->scalar<int64>()() = task.guid;
    }