Method EmitIndexTensor

tensorflow/core/kernels/batch_kernels.cc:644–660

Emits an index tensor, which the Unbatch op will use to un-concatenate the tensor and attribute the pieces to the right batch keys. The index tensor contains, for each input: [batch_key, start_offset, end_offset] where start_offset and end_offset represent the range of entries in the concatenated tensors that belong to that input. Emits the result to the output at 'output_index' using 'context'.

  //
  // Emits the result to the output at 'output_index' using 'context'.
  static Status EmitIndexTensor(OpKernelContext* context, const Batch& batch,
                                int output_index) {
    const TensorShape index_shape({batch.num_tasks(), 3});
    Tensor* index = nullptr;
    TF_RETURN_IF_ERROR(
        context->allocate_output(output_index, index_shape, &index));
    auto index_flat = index->shaped<int64, 2>({batch.num_tasks(), 3});
    size_t offset = 0;
    for (int task_idx = 0; task_idx < batch.num_tasks(); ++task_idx) {
      const BatchTask& task = batch.task(task_idx);
      index_flat(task_idx, 0) = task.guid;
      index_flat(task_idx, 1) = offset;
      index_flat(task_idx, 2) = offset + task.size();
      offset += task.size();
    }
    return Status::OK();
  }

  // Looks up the batcher queue for 'queue_name'. If it did't previously exist,
  // creates it.