tensorflow::Tensor* output = op_context.mutable_output(i);
    const tensorflow::XlaExpression* expr =
        tensorflow::XlaExpression::CastExpressionFromTensor(*output);
    output_values.push_back(expr->AsXlaOp(&xla_builder_));
  }

  absl::Span<const xla::XlaOp> return_values(output_values);
  xla::XlaOp root_value = xla::Tuple(&xla_builder_, return_values);

  TF_ASSIGN_OR_RETURN(XlaComputation computation,