custom_legalization_passes,
                           /*shape_determination_fns=*/{}, arg_shapes,
                           &arg_core_mapping, &per_core_arg_shapes, client);
}

TEST(LegalizeTFTest, RecordsStreamzForSuccessfulLegalizeWithMlirBridge) {
  CellReader<int64_t> compilation_status(kCompilationStatusStreamzName);

  TF_ASSERT_OK_AND_ASSIGN(
      XlaCompiler::CompilationResult result,

        "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
        "@local_xla//xla/pjrt:compile_options_proto_cc",
        "@stablehlo//:register",
    ],
)

tf_cc_test(
    name = "legalize_tf_test",
    srcs = ["legalize_tf_test.cc"],
    deps = [
        ":device_type_proto_cc",
        ":legalize_tf",
        "//tensorflow/compiler/jit",
        "//tensorflow/compiler/mlir/tensorflow",