// CHECK: TPUCompileMlir
  // CHECK: "tf.OpA"(%[[COMPILE_LAUNCH]]#1
  %0:2 = "tf_device.launch"() ({
    %compilation_status, %program = "tf._TPUCompileMlir"() { metadata = "...", mlir_module = "..." } : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
    tf_device.return %compilation_status, %program : tensor<!tf_type.string>, tensor<3x!tf_type.string>

  }

  {
    mutex_lock lock(entry->mu);
    if (compile_state.has_value()) {
      entry->compile_state = *compile_state;
    }
    if (compilation_status.has_value()) {
      entry->compilation_status = *compilation_status;
    }
    if (compilation_result.has_value()) {
      entry->compilation_result = std::move(*compilation_result);
    }
    if (executable.has_value()) {

  CellReader<int64_t> compilation_status(kExportStreamzName);

  TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir"));

  EXPECT_FALSE(ExportFromTensorflowDialectToExecutor(*mlir_module_).ok());

  EXPECT_EQ(compilation_status.Delta(kExportSuccess), 0);
  EXPECT_EQ(compilation_status.Delta(kExportFailed), 1);
}

TEST_F(TensorflowDialectToExecutorTest, PrunesDeadOps) {
  CellReader<int64_t> compilation_status(kExportStreamzName);

  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
  %compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()

  EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldSuccess), 1);
  EXPECT_EQ(compilation_status.Delta(kMlirCombinedOldFailure), 0);
  // Old bridge should never be called at all.
  EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredFailure), 0);
  EXPECT_EQ(compilation_status.Delta(kOldBridgeWithFallbackModeFailure), 0);
  EXPECT_EQ(compilation_status.Delta(kOldBridgeMlirFilteredSuccess), 0);

  EXPECT_EQ(cache_value_1->compilation_status.code(), error::INVALID_ARGUMENT);
  EXPECT_TRUE(cache_value_1->compilation_result != nullptr);
  EXPECT_TRUE(cache_value_1->executable != nullptr);
  EXPECT_EQ(cache_value_1->executable->data, "foo_exe");

  EXPECT_EQ(cache_value_2->compile_state, DeviceCompileState::kCompiling);
  EXPECT_TRUE(cache_value_2->compilation_status.ok());

  CellReader<int64_t> compilation_status(kCompilationStreamz);

  TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));

  TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
      *mlir_module_, /*is_supported_by_replicated_brige*/ true,
      /*is_in_fallback_enabled_mode=*/true));

  EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
                                     mlir::TF::kMlirPh1BridgeCounterV2,

};

TEST_F(SessionClusterTensorflowDialectTest, ClustersTf) {
  CellReader<int64_t> compilation_status(kCompilationStreamz);

  TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));

  TF_EXPECT_OK(
      RunSessionTf2xlaClusteringBridge(*mlir_module_,
                                       /*is_in_fallback_enabled_mode=*/false));
  EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,

  CellReader<int64_t> compilation_status(kCompilationStreamz);

  TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));

  TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));

  EXPECT_EQ(compilation_status.Delta("success"), 1);
}

TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
  CellReader<int64_t> compilation_status(kCompilationStreamz);

    absl::Status compilation_status =
        tensorflow::tf2xla::v1::CompileTensorflowGraphToHlo(
            computation, metadata_proto, use_tuple_args,
            shape_determination_fns, arg_shapes, &arg_core_mapping,
            &per_core_arg_shapes, client, &compilation_result);

    if (!compilation_status.ok()) return compilation_status;

    return compilation_result;
  }
};