def apply_compute_plan(
    model: MLModel,
    compute_plan: MLComputePlan,
    backend_assignment_fn: Optional[
        Callable[[proto.MIL_pb2.Operation, Optional[MLComputePlanDeviceUsage]], Optional[List[str]]]
    ] = None,
) -> MLModel:
    """
    This function takes an MLModel and sets the intended backend attribute
    of each operation in the model as determined by the model's compute plan.

    It updates the 'IntendedBackend' attribute of each operation,
    ensuring the same dispatch even if the model is modified.

    Parameters:
    -----------
    model : MLModel
        The model to which the compute plan will be applied.

    compute_plan : Optional[MLComputePlan]
        The model's compute plan.

    backend_assignment_fn : Callable
        A function that determines the intended backends for each operation.

    Returns:
    --------
    MLModel
        The modified model after applying the backend assignments.
    """

    debug_id = 0
    debug_id_attr_name = "debug_id"

    def get_debug_id(op: proto.MIL_pb2.Operation) -> int:
        attr = op.attributes[debug_id_attr_name]
        values = attr.immediateValue.tensor.strings.values
        if len(values) == 0:
            raise KeyError(
                f"Debug ID attribute '{debug_id_attr_name}' not found in operation: {op.name}"
            )
        value = values[0]
        return int(value)

    def id_generator(op: proto.MIL_pb2.Operation) -> str:
        nonlocal debug_id
        debug_id += 1
        return str(debug_id)

    def backend_assignment_using_preferred_compute_device(
        op: proto.MIL_pb2.Operation, compute_device_usage: Optional[MLComputePlanDeviceUsage]
    ) -> Optional[List[str]]:
        if compute_device_usage is None:
            return None

        preferred_compute_device = compute_device_usage.preferred_compute_device
        return _get_backend_names(compute_device=preferred_compute_device)

    if backend_assignment_fn is None:
        backend_assignment_fn = backend_assignment_using_preferred_compute_device

    model_spec = model.get_spec()

    if model_spec.WhichOneof("Type") != "mlProgram":
        raise ValueError("The apply_compute_plan only supports ML Program.")

    _assign_ids_to_program_operations(
        program_spec=model_spec.mlProgram,
        id_generator=id_generator,
        attr_name=debug_id_attr_name,
    )

    model = MLModel(
        model=model_spec,
        weights_dir=model.weights_dir,
        function_name=model.function_name,
        compute_units=model.compute_unit,
    )

    values = map_model_spec_to_path(
        model_spec=model_spec,
    )

    debug_id_to_path_map = {get_debug_id(op_spec): path for op_spec, path in values}

    values = map_model_structure_to_path(
        model_structure=compute_plan.model_structure,
    )
    path_to_structure_map = {path: structure for structure, path in values}

    def set_intended_backends_for_op(
        op_spec: proto.MIL_pb2.Operation,
    ) -> Optional[List[str]]:
        path = debug_id_to_path_map.get(get_debug_id(op_spec), None)
        op_structure = path_to_structure_map.get(path, None) if path is not None else None
        compute_device_usage = None
        if op_structure is not None:
            compute_device_usage = compute_plan.get_compute_device_usage_for_mlprogram_operation(
                op_structure
            )

        backend_names = backend_assignment_fn(op_spec, compute_device_usage)
        return backend_names

    return set_intended_backends(
        model=model,
        backend_assignment_fn=set_intended_backends_for_op,
    )

    This function takes an MLModel and sets the intended backend attribute
    of each operation in the model as determined by the model's compute plan.

    It updates the 'IntendedBackend' attribute of each operation,
    ensuring the same dispatch even if the model is modified.

    Parameters:
    -----------
    model : MLModel
        The model to which the compute plan will be applied.

    compute_plan : Optional[MLComputePlan]
        The model's compute plan.

    backend_assignment_fn : Callable
        A function that determines the intended backends for each operation.

    Returns:
    --------
    MLModel
        The modified model after applying the backend assignments.
    
coremltools/models/ml_program/experimental/compute_plan_utils.py