def _verify_output_correctness_of_chunks(
    full_model: "_ct.models.MLModel",
    first_chunk_model: _Optional["_ct.models.MLModel"] = None,
    second_chunk_model: _Optional["_ct.models.MLModel"] = None,
    pipeline_model: _Optional["_ct.models.MLModel"] = None,
) -> None:
    """Verifies the end-to-end output correctness of full (original) model versus chunked models"""
    # lazy import avoids circular error
    from coremltools.converters.mil.testing_utils import compute_snr_and_psnr
    from coremltools.converters.mil.testing_utils import (
        random_gen_input_feature_type as random_gen_input_feature_type,
    )

    def report_correctness(original_outputs: _np.ndarray, final_outputs: _np.ndarray, log_prefix: str):
        """ Report PSNR values across two compatible tensors.
        This util is from https://github.com/apple/ml-stable-diffusion/blob/main/python_coreml_stable_diffusion/torch2coreml.py#L80,
        with a slightly modification.
        """
        ABSOLUTE_MIN_PSNR = 35

        _, original_psnr = compute_snr_and_psnr(original_outputs, original_outputs)
        _, final_psnr = compute_snr_and_psnr(original_outputs, final_outputs)

        dB_change = final_psnr - original_psnr
        _logger.info(
            f"{log_prefix}: PSNR changed by {dB_change:.1f} dB ({original_psnr:.1f} -> {final_psnr:.1f})"
        )

        if final_psnr < ABSOLUTE_MIN_PSNR:
            _logger.warning(f"{final_psnr:.1f} dB is low!")
        else:
            _logger.info(
                f"{final_psnr:.1f} dB > {ABSOLUTE_MIN_PSNR} dB (minimum allowed) parity check passed"
            )
        return final_psnr


    # Generate inputs for first chunk and full model
    input_dict = {}
    for input_desc in full_model._spec.description.input:
        input_dict[input_desc.name] = random_gen_input_feature_type(input_desc)

    # Generate outputs for full model
    outputs_from_full_model = full_model.predict(input_dict)

    if pipeline_model is not None:
        outputs_from_pipeline_model = pipeline_model.predict(input_dict)
        final_outputs = outputs_from_pipeline_model

    elif first_chunk_model is not None and second_chunk_model is not None:
        # Generate outputs for first chunk
        outputs_from_first_chunk_model = first_chunk_model.predict(input_dict)

        # Prepare inputs for second chunk model from first chunk's outputs and regular inputs
        second_chunk_input_dict = {}
        for input_desc in second_chunk_model._spec.description.input:
            if input_desc.name in outputs_from_first_chunk_model:
                second_chunk_input_dict[input_desc.name] = outputs_from_first_chunk_model[
                    input_desc.name
                ]
            else:
                second_chunk_input_dict[input_desc.name] = input_dict[input_desc.name]

        # Generate output for second chunk model
        outputs_from_second_chunk_model = second_chunk_model.predict(second_chunk_input_dict)
        final_outputs = outputs_from_second_chunk_model
    else:
        raise ValueError("Either a single Pipeline model or two model chunks should be provided.")

    # Verify correctness across all outputs from second chunk and full model
    for out_name in outputs_from_full_model.keys():
        report_correctness(
            original_outputs=outputs_from_full_model[out_name],
            final_outputs=final_outputs[out_name],
            log_prefix=f"{out_name}",
        )
Verifies the end-to-end output correctness of full (original) model versus chunked models
coremltools/models/utils.py