const XlaCompiler::CompilationResult* compilation_result = nullptr;
  xla::LocalExecutable* xla_executable = nullptr;

  TF_EXPECT_OK(xla_device_compiler->CompileIfNeeded(
      options, fn, args, XlaCompiler::CompileOptions{},
      DeviceCompileMode::kStrict, profiler_, &compilation_result,
      &xla_executable));

  EXPECT_TRUE(compilation_result != nullptr);
  EXPECT_TRUE(xla_executable != nullptr);

}

void RunInThreadPoolIfCollectivesPresent(
    const XlaCompiler::CompilationResult& compilation_result,
    std::function<void()> execution_fn) {
  // If we are using collectives, we need to run in a separate threadpool.
  if (compilation_result.collective_info.has_value()) {
    GetOrCreateThreadPoolForCollective(*compilation_result.collective_info)
        ->Schedule(execution_fn);
  } else {

    OpKernelContext* ctx,
    const XlaCompiler::CompilationResult& compilation_result,
    int missing_ctx_input_prefix) {
  std::vector<VariableInfo> out;
  out.reserve(compilation_result.resource_updates.size());
  for (int i = 0; i < compilation_result.resource_updates.size(); ++i) {
    const XlaCompiler::ResourceUpdate& write =
        compilation_result.resource_updates[i];

    const XlaCompiler::CompilationResult* compilation_result, int output_num);

// Converts input tensors and variables which are parameters of the
// XlaComputation into PjRtBuffers to be fed as input to the
// PjRtLoadedExecutable.
//
// Assumes that the first `num_missing_prefix_ctx_inputs` inputs to the
// compilation_result are missing in `inputs` and adjusts indexing into `inputs`
// accordingly.

        return %cst : tensor<2xi32>
    }
  })";

  auto compilation_result = CompileMlirModule(
      kHasReturnValuesAndNoMetadataRetvals,
      ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED);
  EXPECT_TRUE(compilation_result.ok());

  // Ensure that the compilation result contains a constant.
  EXPECT_THAT(compilation_result,
              ComputationProtoContains("opcode:.*constant"));
}

        return %0 : tensor<32x64xf32>
    }
  })";

  auto compilation_result =
      CompileWithComputation(CreateTestMlirToHloArgs(kHasReturnValues));
  EXPECT_TRUE(compilation_result.ok());

  // The StatelessRandomNormal must not be replaced by a literal tensor.
  EXPECT_THAT(compilation_result,
              ComputationProtoContains("tf.StatelessRandomNormal"));
}

  GuaranteedConsts consts;

  *compilation_result = {};

  TF_RETURN_IF_ERROR(CompileTFFunctionToHlo(
      *flib_def, versions.producer(), shape_determination_funcs, arg_shapes,
      consts, func, metadata, client, arg_core_mapping, per_core_arg_shapes,
      use_tuple_args, compilation_result));

  return PopulateInputOutputAliasing(main_fn, compilation_result,
                                     use_tuple_args);

    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    XlaCompilationResult* compilation_result);

// Runs MLIR Bridge on a MLIR module.
//
// If lower_to_xla_hlo is true then compiles down into XLA HLO, generates all
// accompanying metadata and stores them in CompilationResult.
//
// If enable_op_fallback is set to false, graph is legalized only if the graph

    uint64 signature_hash, const XlaCompiler::Options& options,
    const XlaCompiler::CompilationResult& compilation_result,
    const ExecutableType& executable,
    DeviceCompilerClient<ExecutableType, ClientType>* compiler_client) const {
  XlaSerializedCacheEntry serialized_entry;
  const xla::HloModuleProto& hlo_module =
      compilation_result.computation->proto();
  *serialized_entry.mutable_key() =

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(
      compilation_result.computation->proto(), module_config);
  if (!status_or_hlo_module.ok()) {