Method RegisterInput

tensorflow/core/kernels/batch_kernels.cc:244–281

Ingests data from one invocation of the batch op. The data is enqueued to be combined with others into a batch, asynchronously.

  // Ingests data from one invocation of the batch op. The data is enqueued to
  // be combined with others into a batch, asynchronously.
  Status RegisterInput(int64 guid, OpKernelContext* context,
                       const string& batcher_queue_name,
                       AsyncOpKernel::DoneCallback done_callback) {
    std::unique_ptr<BatchTask> batch_components(new BatchTask);
    batch_components->guid = guid;
    batch_components->propagated_context = Context(ContextKind::kThread);
    OpInputList tensors;
    TF_RETURN_IF_ERROR(context->input_list("in_tensors", &tensors));
    for (int i = 0; i < tensors.size(); ++i) {
      const Tensor& tensor = tensors[i];
      if (tensor.shape().dims() == 0) {
        return errors::InvalidArgument(
            "Batching input tensors must have at least one dimension");
      }
      if (tensors.size() >= 2 &&
          tensor.shape().dim_size(0) != tensors[0].shape().dim_size(0)) {
        return errors::InvalidArgument(
            "Batching input tensors supplied in a given op invocation must "
            "have equal 0th-dimension size");
      }
      batch_components->inputs.push_back(tensor);
    }
    OpInputList captured_tensors;
    const auto captured_status =
        context->input_list("captured_tensors", &captured_tensors);
    if (captured_status.ok()) {
      for (const Tensor& captured_tensor : captured_tensors) {
        batch_components->captured_inputs.push_back(captured_tensor);
      }
    }
    batch_components->context = context;
    batch_components->done_callback = std::move(done_callback);

    BatcherQueue* batcher_queue;
    TF_RETURN_IF_ERROR(
        LookupOrCreateBatcherQueue(batcher_queue_name, &batcher_queue));
    return batcher_queue->Schedule(&batch_components);
  }

 private:
  BatchResource() = default;