*out_context = device_context;
  return absl::OkStatus();
}

// Warn about XLA_CPU/XLA_GPU exactly once.
static void ShowXlaDeviceDeprecationWarning(
    absl::string_view compilation_device_name) {
  static absl::once_flag once;
  if (absl::StrContains(compilation_device_name, "CPU") ||
      absl::StrContains(compilation_device_name, "GPU")) {
    absl::call_once(once, [] {

#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h"

#include <cstdint>
#include <memory>
#include <string>
#include <variant>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

    }
  }

  if (!is_module_updated) {
    VLOG(2) << "MLIR module is not updated. Using the original graph. "
            << "Do not convert mlir module back to graph";
    return absl::OkStatus();
  }
  GraphExportConfig export_config;
  absl::flat_hash_set<Node*> control_ret_nodes;

  timings.Reset({kTfMlirCategory, "convert_mlir_to_graph"});
  // Some or all passes are enabled. Convert MLIR module and return back

#include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.h"

#include <cstdint>
#include <optional>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/Support/MathExtras.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project

    TF_DeleteStatus(status);
    return ok;
  };
  StreamExecutor* executor = GetExecutor(0);
  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
  std::function<absl::Status()> callback = []() -> absl::Status {
    return absl::OkStatus();
  };
  TF_ASSERT_OK(stream->DoHostCallbackWithStatus(callback));
}

TEST_F(StreamExecutorTest, HostCallbackError) {

#include "tsl/platform/errors.h"

namespace tensorflow {
namespace {
constexpr absl::string_view kGroupSizeAttrName =
    "tf2xla.collective_info.group_size";
constexpr absl::string_view kGroupKeyAttrName =
    "tf2xla.collective_info.group_key";

// Extracts shape from XlaArgument as TensorShape. If shape is a xla::Shape,
// that is converted to a TensorShape.
absl::StatusOr<TensorShape> GetTensorShapeFromXlaArgument(
    const XlaArgument& arg) {

  } else {
    status = host_to_device_stream_->DoHostCallback([ref, done]() {
      ref.Unref();
      done(absl::OkStatus());
    });
    if (!status.ok()) {
      done(status);
    }
  }
}

void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor,
                                             absl::string_view tensor_name,
                                             Device* device, Tensor* cpu_tensor,

==============================================================================*/

#include "tensorflow/compiler/jit/shape_inference.h"

#include <cstdint>
#include <map>
#include <vector>

#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/shape_inference_helpers.h"
#include "tensorflow/core/common_runtime/shape_refiner.h"
#include "tensorflow/core/framework/function.h"

#include <Python.h>

#include <fstream>
#include <memory>
#include <optional>
#include <string>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
#include "tensorflow/c/kernels.h"
#include "tensorflow/c/tf_status.h"

        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "//tensorflow/compiler/mlir/quantization/common/quantization_lib",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Pass",