se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =

    Graph* g, Node* n, FunctionLibraryDefinition* fld) {
  TF_RET_CHECK(n->IsWhileNode());

  // Check if there is any lifted args in body function.
  NameAttrList body_func;
  TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "body", &body_func));
  const FunctionDef* body_function_def = fld->Find(body_func.name());
  TF_RET_CHECK(body_function_def);

  if (!HasLiftedArgs(*body_function_def)) {

                             mlir::MLIRContext* mlir_context,
                             mlir::OwningOpRef<mlir::ModuleOp>* mlir_module) {
  TF_RET_CHECK(!serialized_mlir_module.empty())
      << "unexpected empty serialized MLIR module string";
  TF_RET_CHECK(mlir_module) << "unexpected null MLIR module pointer";

  // Make sure we catch any error reported by MLIR and forward it to the TF
  // error reporting system.

      options, result, client_->default_device_ordinal());
  TF_ASSIGN_OR_RETURN(
      auto executables,
      client_->Compile(*result.computation, argument_layouts, build_options));
  TF_RET_CHECK(executables.size() == 1);
  return std::move(executables[0]);
}

absl::StatusOr<std::string> XlaDeviceCompilerClient::SerializeExecutable(
    const xla::LocalExecutable& executable) {

    TF_RETURN_IF_ERROR(
        ctx->allocate_temp(output_dtype, output_shape, &output_tensor));
    if (output_tensor.TotalBytes() > 0) {
      XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor);
      TF_RET_CHECK(xla_tensor);
      xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num}));
      if (use_multiple_streams) {
        xla_tensor->ResetDefinitionEvent(definition_event, stream);
      }
    }

      // Handles tf.resource variables.
      TF_RET_CHECK(input->dtype() == DT_RESOURCE);
      const VariableInfo& variable = *variable_info_lookup[input_num];
      arg.kind = XlaCompiler::Argument::kResource;
      arg.resource_kind = XlaResource::kVariable;
      arg.definition_stack_trace = variable.definition_stack_trace();
      TF_RET_CHECK(variable.var() && variable.var()->is_initialized);

      if (le->IsControlEdge()) {
        control_outputs.insert(le->dst());
      } else {
        TF_RET_CHECK(le->src_output() < num_outputs);
        Node* output_node = le->dst();

        if (add_edges_to_output_of_downstream_nodes) {
          TF_RET_CHECK(output_node->type_string() == kXlaClusterOutput)
              << le->DebugString();
          nodes_to_remove.push_back(output_node);

  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";

  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("inputs")))
        << "inputs missing in entry function attribute";
    TF_RET_CHECK(mlir::isa<mlir::StringAttr>(dict_attr.get("outputs")))
        << "outputs missing in entry function attribute";

    const Graph& graph, const GraphShapeInfo& shape_info,
    std::map<string, std::vector<PartialTensorShape>> expected_shapes) {
  for (Node* node : graph.op_nodes()) {
    auto sit = shape_info.find(node->name());
    TF_RET_CHECK(sit != shape_info.end())
        << "Missing shape information for node " << node->name();
    std::vector<PartialTensorShape> shapes;
    for (const auto& output : sit->second) shapes.push_back(output.shape);