def bisect_model(
    model: _Union[str, "_ct.models.MLModel"],
    output_dir: str,
    merge_chunks_to_pipeline: _Optional[bool] = False,
    check_output_correctness: _Optional[bool] = True,
):
    """
    Utility function to split a mlpackage model into two mlpackages of approximately same file size.

    Parameters
    ----------
    model: str or MLModel
        Path to the mlpackage file, or a Core ML model, to be split into two mlpackages of approximately same file size.

    output_dir: str
        Path to output directory where the two model chunks / pipeline model would be saved.

        If the `model` is `{path}/{model_name}.mlpackage`, the chunk models are going to be saved as:
        1. first chunk model: `{output_dir}/{model_name}_chunk1.mlpackage`
        2. second chunk model: `{output_dir}/{model_name}_chunk2.mlpackage`
        3. chunked pipeline model: `{output_dir}/{model_name}_chunked_pipeline.mlpackage`

        If the `model` is type of `MLModel`, the chunk models are saved as:
        1. first chunk model: `{output_dir}/chunk1.mlpackage`
        2. second chunk model: `{output_dir}/chunk2.mlpackage`
        3. chunked pipeline model: `{output_dir}/chunked_pipeline.mlpackage`

    merge_chunks_to_pipeline: bool
        If True, model chunks are managed inside a single pipeline model for easier asset maintenance.

    check_output_correctness: bool
        - If True, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB.
        - Enabling this feature uses more memory. Disable it if your machine runs out of memory.

    Examples
    --------
    .. sourcecode:: python

        import coremltools as ct

        model_path = "my_model.mlpackage"
        output_dir = "./output/"

        # The following code will produce two smaller models:
        # `./output/my_model_chunk1.mlpackage` and `./output/my_model_chunk2.mlpackage`
        # It also compares the output numerical of the original Core ML model with the chunked models.
        ct.models.utils.bisect_model(
            model_path,
            output_dir,
        )

        # The following code will produce a single pipeline model `./output/my_model_chunked_pipeline.mlpackage`
        ct.models.utils.bisect_model(
            model_path,
            output_dir,
            merge_chunks_to_pipeline=True,
        )

        # You can also pass the MLModel object directly
        mlmodel = ct.models.MLModel(model_path)
        ct.models.utils.bisect_model(
            mlmodel,
            output_dir,
            merge_chunks_to_pipeline=True,
        )
    """
    def get_pymil_prog_and_spec_from_model(model):

        # get the model spec and weight directory
        if isinstance(model, str):
            spec = load_spec(model)
            weights_dir = _try_get_weights_dir_path(model)
        else:
            spec = model._spec
            weights_dir = model.weights_dir

        # convert the model spec into pymil program,
        # we also convert operations into type of List
        prog = _milproto_to_pymil.load(
            spec,
            spec.specificationVersion,
            weights_dir,
        )
        if len(prog.functions) > 1 or "main" not in prog.functions:
            raise ValueError("'bisect_model' only support model with a single 'main' function.")

        func = prog.functions["main"]
        func.operations = list(func.operations)

        return prog, spec

    # check the input type of model
    if not isinstance(model, (str, _ct.models.MLModel)):
        raise ValueError(f"'model' must be type of [str, MLModel]. Got {type(model)}.")

    # The below implementation assumes that the model is single function, with a "main" function.
    prog, spec = get_pymil_prog_and_spec_from_model(model)
    spec_version = spec.specificationVersion

    # Compute the incision point by bisecting the program based on weights size
    op_idx, first_chunk_weights_size, total_weights_size = _get_op_idx_split_location(prog)
    main_block = prog.functions["main"]
    incision_op = main_block.operations[op_idx]
    _logger.info(
        f"The incision op: name={incision_op.name}, type={incision_op.op_type}, index={op_idx}/{len(main_block.operations)}"
    )
    _logger.info(f"First chunk size = {first_chunk_weights_size:.2f} MB")
    _logger.info(f"Second chunk size = {total_weights_size - first_chunk_weights_size:.2f} MB")

    # Build first chunk (in-place modifies prog by declaring early exits and removing unused subgraph)
    prog_chunk1 = _make_first_chunk_prog(prog, op_idx)

    # Build the second chunk
    # when the first chunk is created, the prog is modified in-place, so we need to re-convert a new pymil
    # program for the second chunk.
    prog_chunk2 = _make_second_chunk_prog(
        get_pymil_prog_and_spec_from_model(model)[0],
        op_idx,
    )

    # Convert the MIL Program objects into MLModels
    # We skip_model_load if check_output_correctness=False
    _logger.info("Converting the two programs")
    model_chunk1 = _mil_convert(
        prog_chunk1,
        convert_to="mlprogram",
        convert_from="milinternal",
        specification_version=spec_version,
        compute_units=_ct.ComputeUnit.CPU_ONLY,
        skip_model_load=(not check_output_correctness),
    )
    del prog_chunk1
    _gc.collect()
    _logger.info("Conversion of first chunk done.")

    model_chunk2 = _mil_convert(
        prog_chunk2,
        convert_to="mlprogram",
        convert_from="milinternal",
        specification_version=spec_version,
        compute_units=_ct.ComputeUnit.CPU_ONLY,
        skip_model_load=(not check_output_correctness),
    )
    del prog_chunk2
    _gc.collect()
    _logger.info("Conversion of second chunk done.")

    # Verify output correctness
    if check_output_correctness:
        _logger.info("Verifying output correctness of chunks")

        if isinstance(model, str):
            mlmodel = _ct.models.MLModel(model, compute_units=_ct.ComputeUnit.CPU_ONLY)
        else:
            mlmodel = model

        _verify_output_correctness_of_chunks(
            full_model=mlmodel,
            first_chunk_model=model_chunk1,
            second_chunk_model=model_chunk2,
        )

    # save model chunks
    _os.makedirs(output_dir, exist_ok=True)

    if isinstance(model, str):
        mlpackage_name = _os.path.basename(model)
        name, _ = _os.path.splitext(mlpackage_name)
        name += "_"
    else:
        name = ""

    if merge_chunks_to_pipeline:
        # Make a single pipeline model to manage the model chunks
        pipeline_model = make_pipeline(model_chunk1, model_chunk2)
        out_path_pipeline = _os.path.join(output_dir, name + "chunked_pipeline.mlpackage")
        pipeline_model.save(out_path_pipeline)

        # reload to ensure CPU placement
        if check_output_correctness:
            _logger.info("Verifying output correctness of pipeline model")
            pipeline_model = _ct.models.MLModel(
                out_path_pipeline, compute_units=_ct.ComputeUnit.CPU_ONLY
            )
            _verify_output_correctness_of_chunks(
                full_model=mlmodel,
                pipeline_model=pipeline_model,
            )
    else:
        # Save the chunked models to disk
        out_path_chunk1 = _os.path.join(output_dir, name + "chunk1.mlpackage")
        out_path_chunk2 = _os.path.join(output_dir, name + "chunk2.mlpackage")
        model_chunk1.save(out_path_chunk1)
        model_chunk2.save(out_path_chunk2)
        _logger.info(
            f"Saved chunks in {output_dir} with the suffix _chunk1.mlpackage and _chunk2.mlpackage"
        )

    Utility function to split a mlpackage model into two mlpackages of approximately same file size.

    Parameters
    ----------
    model: str or MLModel
        Path to the mlpackage file, or a Core ML model, to be split into two mlpackages of approximately same file size.

    output_dir: str
        Path to output directory where the two model chunks / pipeline model would be saved.

        If the `model` is `{path}/{model_name}.mlpackage`, the chunk models are going to be saved as:
        1. first chunk model: `{output_dir}/{model_name}_chunk1.mlpackage`
        2. second chunk model: `{output_dir}/{model_name}_chunk2.mlpackage`
        3. chunked pipeline model: `{output_dir}/{model_name}_chunked_pipeline.mlpackage`

        If the `model` is type of `MLModel`, the chunk models are saved as:
        1. first chunk model: `{output_dir}/chunk1.mlpackage`
        2. second chunk model: `{output_dir}/chunk2.mlpackage`
        3. chunked pipeline model: `{output_dir}/chunked_pipeline.mlpackage`

    merge_chunks_to_pipeline: bool
        If True, model chunks are managed inside a single pipeline model for easier asset maintenance.

    check_output_correctness: bool
        - If True, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB.
        - Enabling this feature uses more memory. Disable it if your machine runs out of memory.

    Examples
    --------
    .. sourcecode:: python

        import coremltools as ct

        model_path = "my_model.mlpackage"
        output_dir = "./output/"

        # The following code will produce two smaller models:
        # `./output/my_model_chunk1.mlpackage` and `./output/my_model_chunk2.mlpackage`
        # It also compares the output numerical of the original Core ML model with the chunked models.
        ct.models.utils.bisect_model(
            model_path,
            output_dir,
        )

        # The following code will produce a single pipeline model `./output/my_model_chunked_pipeline.mlpackage`
        ct.models.utils.bisect_model(
            model_path,
            output_dir,
            merge_chunks_to_pipeline=True,
        )

        # You can also pass the MLModel object directly
        mlmodel = ct.models.MLModel(model_path)
        ct.models.utils.bisect_model(
            mlmodel,
            output_dir,
            merge_chunks_to_pipeline=True,
        )
    
coremltools/models/utils.py