llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
                          custom_legalization_passes);

// Apply shape, description, and resource information to inputs and outputs
// in the XlaCompilationResult. This should be called after
// compilation_result->computation was set.
ABSL_DEPRECATED("Not meant to be used directly and should be a util.")
Status PopulateResultIOInfo(

    llvm::cl::init(false));
}  // namespace

namespace tensorflow {

namespace {

mlir::LogicalResult PrintHloModuleText(
    const XlaCompilationResult& compilation_result, llvm::raw_ostream& output) {
  const xla::HloModuleConfig module_config(
      compilation_result.computation->GetProgramShape().value());
  auto status_or_hlo_module = xla::HloModule::CreateFromProto(

  return mlir_to_hlo_args;
}

class CompileTFGraphTest : public ::testing::Test {
 public:
  absl::StatusOr<XlaCompilationResult> CompileWithComputation(
      const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>
          computation) {
    XlaCompilationResult compilation_result;

    se::Platform* platform =
        se::PlatformManager::PlatformWithName(kPlatformName).value();
    auto client =

// compilation (HLO) using the XlaCompiler and compiles the resulting
// XlaCompilationResult into an `ExecutableType` (eg. xla::LocalExecutable) by
// calling `ClientType` (eg. xla::LocalClient).
//
// Caches the compiled XlaCompilationResult and Executable using a
// DeviceCompilationCache. Compilation is done only when there's a cache miss.
//

  if (VLOG_IS_ON(2))
    tensorflow::DumpMlirOpToFile("build_hlo_tf_after", module_op);

  return absl::OkStatus();
}

Status PopulateCollectiveInfo(mlir::ModuleOp module_op,
                              XlaCompilationResult* compilation_result) {
  // The StringRef cast is necessary before cxx14.
  mlir::IntegerAttr group_key_attr =
      module_op->getAttrOfType<mlir::IntegerAttr>(

// Get-or-create thread pool for a given collective.
static thread::ThreadPool* GetOrCreateThreadPoolForCollective(
    const XlaCompilationResult::CollectiveInfo& collective_info) {
  static absl::Mutex m(absl::kConstInit);
  static auto& thread_pool_cache ABSL_GUARDED_BY(m) =
      *new absl::node_hash_map<XlaCompilationResult::CollectiveInfo,
                               thread::ThreadPool>();
  absl::MutexLock l(&m);