def _convert_and_retrieve_jit_module_mapping(
    model: ScriptModule,
    annotator: Optional[TorchScriptModuleAnnotator] = None,
    **converter_kwargs,
) -> Tuple[MLModel, Dict[TorchScriptModuleInfo.Key, TorchScriptModuleMappingInfo]]:
    """
    Converts a TorchScript model to a Core ML model and returns the
    mapping information mapping between the source TorchScript operations and their
    corresponding operations in the converted Core ML model.

    This function takes a TorchScript model, performs  inlining and annotation,
    converts it to a Core ML model using the specified conversion
    parameters, and constructs a mapping that associates nodes in the source
    module with their corresponding operations in the Core ML model.

    Parameters
    ----------
    model : ScriptModule
        The source model.

    annotator : Optional[TorchScriptModuleAnnotator]
        The TorchScript module annotator, if the module has been annotated.

    **converter_kwargs:
        Additional keyword arguments to be passed to the Core ML conversion process.
        These can include options for optimization, input/output specifications, etc.

    Returns
    -------
    Tuple[MLModel, Dict[TorchScriptModuleInfo.Key, TorchScriptModuleMappingInfo]]
        A tuple containing:
            - The converted Core ML model.
            -  A dictionary mapping keys from `TorchScriptModuleInfo` to their corresponding
               ``TorchModuleMLModelMappingInfo`` instances. This mapping provides detailed
               information about how each node in the source module corresponds to operations
               in the target Core ML model.

    Examples
    --------
    .. sourcecode:: python
    
        (
            target_model,
            module_mapping,
        ) = coremltools.models.ml_program.experimental.torch._convert_and_retrieve_jit_module_mapping(
            model=traced_model,
            inputs=[
                coremltools.TensorType(
                    name="x", shape=example_inputs[0].shape, dtype=np.float16
                ),
                coremltools.TensorType(
                    name="y", shape=example_inputs[1].shape, dtype=np.float16
                ),
            ],
            minimum_deployment_target=coremltools.target.iOS16,
            compute_units=coremltools.ComputeUnit.CPU_ONLY,
            skip_model_load=True,
        )
    """
    # Importing here to avoid circular import issues
    from .....converters import convert

    if annotator is None:
        annotator = inline_and_annotate_module(model=model)

    target_model = convert(
        model=model,
        **converter_kwargs,
    )

    op_translation_map = _retrieve_source_to_target_op_map(
        target_model=target_model,
        function_name="main",
    )

    module_call_stack = annotator.module_call_stack
    module_map_infos = OrderedDict()

    def is_module_output(
        op: proto.MIL_pb2.Operation, output_names: List[str]
    ) -> List[proto.MIL_pb2.Operation]:
        for output in op.outputs:
            name = output.name
            for output_name in output_names:
                if name == output_name or name.startswith(f"{output_name}_cast"):
                    return True

        return False

    def _get_module_output_names(
        module_info: TorchScriptModuleInfo,
    ) -> List[str]:
        result = []
        node_infos = annotator.node_infos
        for output_name in module_info.output_names:
            node_info = node_infos.get(output_name, None)
            if node_info.kind == "prim::TupleConstruct":
                result.extend(node_info.input_names)
            else:
                result.append(output_name)

        return result

    for node_info in annotator.node_infos.values():
        module_ops = [
            op_translation_map.output_name_to_target_ops.get(name, [])
            for name in node_info.output_names
        ]
        module_ops = sum(module_ops, [])
        for module_key in node_info.modules:
            module_info = module_call_stack[module_key]
            module_map_info = module_map_infos.setdefault(
                module_key,
                TorchScriptModuleMappingInfo(
                    source=module_info,
                    source_to_target_ops_mapping=OrderedDict(),
                    deps=[],
                    outputs=[],
                    submodules=[],
                ),
            )

            module_output_names = _get_module_output_names(module_info)
            if len(module_map_info.outputs) < len(module_output_names):
                for op in module_ops:
                    if is_module_output(op, output_names=module_output_names):
                        module_map_info.outputs.append(op)

            ops = module_map_info.source_to_target_ops_mapping.get(node_info, [])
            ops.extend(module_ops)
            module_map_info.source_to_target_ops_mapping[node_info] = ops

    # Set module deps and submodules
    graph = _build_graph(nodes=annotator.node_infos)
    for module_key, module_map_info in module_map_infos.items():
        submodules = [module_call_stack[key] for key in module_map_info.source.submodules]
        module_map_info.submodules = submodules
        deps = _get_module_deps(
            parent_module=module_map_info.source,
            submodules=submodules,
            nodes=annotator.node_infos,
            graph=graph,
        )
        module_map_info.deps = {
            key: [module_map_infos[value.key] for value in values] for key, values in deps.items()
        }

    return (target_model, module_map_infos)

    Converts a TorchScript model to a Core ML model and returns the
    mapping information mapping between the source TorchScript operations and their
    corresponding operations in the converted Core ML model.

    This function takes a TorchScript model, performs  inlining and annotation,
    converts it to a Core ML model using the specified conversion
    parameters, and constructs a mapping that associates nodes in the source
    module with their corresponding operations in the Core ML model.

    Parameters
    ----------
    model : ScriptModule
        The source model.

    annotator : Optional[TorchScriptModuleAnnotator]
        The TorchScript module annotator, if the module has been annotated.

    **converter_kwargs:
        Additional keyword arguments to be passed to the Core ML conversion process.
        These can include options for optimization, input/output specifications, etc.

    Returns
    -------
    Tuple[MLModel, Dict[TorchScriptModuleInfo.Key, TorchScriptModuleMappingInfo]]
        A tuple containing:
            - The converted Core ML model.
            -  A dictionary mapping keys from `TorchScriptModuleInfo` to their corresponding
               ``TorchModuleMLModelMappingInfo`` instances. This mapping provides detailed
               information about how each node in the source module corresponds to operations
               in the target Core ML model.

    Examples
    --------
    .. sourcecode:: python
    
        (
            target_model,
            module_mapping,
        ) = coremltools.models.ml_program.experimental.torch._convert_and_retrieve_jit_module_mapping(
            model=traced_model,
            inputs=[
                coremltools.TensorType(
                    name="x", shape=example_inputs[0].shape, dtype=np.float16
                ),
                coremltools.TensorType(
                    name="y", shape=example_inputs[1].shape, dtype=np.float16
                ),
            ],
            minimum_deployment_target=coremltools.target.iOS16,
            compute_units=coremltools.ComputeUnit.CPU_ONLY,
            skip_model_load=True,
        )
    
coremltools/models/ml_program/experimental/torch/debugging_utils.py