return LegalizeWithMlirBridge(
      mlir_to_hlo_args, metadata_proto, use_tuple_args,
      /*device_type=*/"XLA_TPU_JIT",
      /*shape_determination_fns=*/{}, arg_shapes, &arg_core_mapping,
      &per_core_arg_shapes, custom_legalization_passes,
      compilation_result.get());
}

TEST(LegalizeWithMlirBridge, LegalizesToMhloProto) {
  auto result = LegalizeMlirModule(kMlirModuleStr);

// Compiles a serialized MLIR module into XLA HLO, generates all accompanying
// metadata and stores them in CompilationResult.
absl::StatusOr<XlaCompilationResult> LegalizeWithMlirBridge(
    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,

                           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,