auto compilation_result = std::make_unique<XlaCompilationResult>();

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

absl::StatusOr<XlaCompiler::CompilationResult> LegalizeMlirModule(

    const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    const std::vector<tensorflow::TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,

    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    llvm::StringRef device_type,
    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    const std::vector<tensorflow::TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,

    options.use_multiple_streams = true;
    options.allowed_devices = gpu_ids;
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{
        UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
    options.shape_determination_fns = {shape_representation_fns};
    auto device = std::make_unique<XlaDevice>(session_options, options);

  options.device_ordinal = 0;
  options.compilation_device_name = DEVICE_CPU_XLA_JIT;
  options.use_multiple_streams = false;
  XlaShapeLayoutHelpers::ShapeDeterminationFns shape_representation_fns{
      UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()};
  options.shape_determination_fns = {shape_representation_fns};
  auto device = std::make_unique<XlaDevice>(session_options, options);

                   legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
                   /*use_tuple_args=*/true, /*enable_op_fallback=*/false,
                   /*shape_determination_fns=*/{}, &compilation_result)
                   .status());

  EXPECT_FALSE(GetPassThatChangedIdentity().empty());
}

}  // namespace

  return LegalizeTfToHlo(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, client,
                         compilation_result.get());
}

                     mlir_to_hlo_args, metadata_proto, use_tuple_args,
                     /*device_type=*/"XLA_TPU_JIT", custom_legalization_passes,
                     /*shape_determination_fns=*/{}, arg_shapes,
                     &arg_core_mapping, &per_core_arg_shapes, client)
                     .status());
  }
};

TEST_F(LegalizeTFQuantTest, LegalizesModuleWithTFUniformQuantization) {

absl::Status CompileTensorflowGraphToHlo(
    const std::variant<tpu::MlirToHloArgs, tpu::FunctionToHloArgs>& computation,
    const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_funcs,
    const std::vector<tensorflow::TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes,