Method SetOutputExpression

tensorflow/compiler/tf2xla/xla_op_kernel.cc:463–495

}

void XlaOpKernelContext::SetOutputExpression(int index,
                                             const XlaExpression& expression) {
  Status status = [&] {
    // The step's default allocator is the dummy XlaCompilationAllocator which
    // simply allocates a metadata buffer to hold the expression to which it
    // corresponds.
    // Provides a special behavior for DT_VARIANT and other types that are not
    // trivially copyable. In those cases, allocate a tensor of type DT_UINT8.
    if (!DataTypeCanUseMemcpy(expression.dtype())) {
      // tensor_data() is not supported for tensors that cannot be copied via
      // memcpy, as the copy logic might try to inspect the stored data (e.g.
      // a std::string). This is likely to fail, as the data is invalid given
      // that it actually encodes an XlaExpression. Using a uint8 tensor is
      // always safe, so simply do that.
      // TODO(jpienaar): This should be refactored to stop masquerading
      // XlaExpressions as Tensors.
      Tensor output;
      TensorShape tensor_shape;
      TF_RETURN_IF_ERROR(
          context_->allocate_temp(DT_UINT8, tensor_shape, &output));
      context_->set_output(index, output);
    } else {
      Tensor* output = nullptr;
      TF_ASSIGN_OR_RETURN(TensorShape shape, expression.GetShape());
      TF_RETURN_IF_ERROR(context_->allocate_output(index, shape, &output));
    }
    AssignExpressionToTensor(context_->mutable_output(index), expression);
    return Status::OK();
  }();
  if (!status.ok()) {
    SetStatus(status);
  }
}

xla::PrimitiveType XlaOpKernelContext::output_xla_type(int index) {
  xla::PrimitiveType type;