Method GetElementComponentFromBatch

tensorflow/core/kernels/priority_queue.cc:117–128

static */

/* static */
Status PriorityQueue::GetElementComponentFromBatch(
    const PriorityQueue::Tuple& tuple, int index, int component,
    OpKernelContext* ctx, PersistentTensor* out_tensor) {
  TensorShape element_shape(tuple[component].shape());
  element_shape.RemoveDim(0);
  Tensor* element_access = nullptr;
  TF_RETURN_IF_ERROR(ctx->allocate_persistent(
      tuple[component].dtype(), element_shape, out_tensor, &element_access));
  TF_RETURN_IF_ERROR(
      batch_util::CopySliceToElement(tuple[component], element_access, index));
  return Status::OK();
}

void PriorityQueue::TryEnqueueMany(const Tuple& tuple, OpKernelContext* ctx,
                                   DoneCallback callback) {