CellReader<int64_t> reader(kMetric);

  ASSERT_THAT(failed(), Not(IncrementedOrFiltered(reader.Delta(kOkStatus), 1)));
}

tensorflow::XlaCompilationResult CreateXlaComputationResult(
    const char* hlo_name) {
  auto result = tensorflow::XlaCompilationResult();
  xla::HloModuleProto hlo;
  hlo.set_name(hlo_name);
  result.computation = std::make_shared<xla::XlaComputation>(hlo);
  return result;
}

                           arg_core_mapping, per_core_arg_shapes);
  if (!sharding_result.ok()) {
    return sharding_result;
  }
  return compiled_mlir;
}

absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

//  arg_core_mapping - Which args go on which cores.
//  per_core_arg_shapes - For each core, the shapes for each argument.
//  client - The Xla Compilation client.
absl::StatusOr<tensorflow::XlaCompilationResult> LegalizeMlirToHlo(
    const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,

    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.
//

    }
  })";
  SetPassThatChangedIdentity("");
  std::vector<::tensorflow::TensorShape> arg_shapes = {{1}};
  auto compilation_result = tensorflow::XlaCompilationResult();

  TF_EXPECT_OK(tensorflow::CompileSerializedMlirToXlaHlo(
                   legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
                   /*use_tuple_args=*/true, /*enable_op_fallback=*/false,

  std::vector<ShardingAndIndex> arg_core_mapping;
  std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
  auto compilation_result = std::make_unique<XlaCompilationResult>();

  return LegalizeTfToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
                         /*device_type=*/"XLA_TPU_JIT",
                         /*shape_determination_fns=*/{}, arg_shapes,

  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);