def set_intended_backends(
    model: MLModel,
    backend_assignment_fn: Callable[[proto.MIL_pb2.Operation], Optional[List[str]]],
) -> MLModel:
    """
    Assigns intended backends to operations in the given model.

    This function creates a new MLModel with updated backend assignments for each operation.
    It traverses the entire model structure, applying the provided backend assignment function
    to determine the intended backends for each operation.

    Parameters
    ----------
    model : MLModel
        The input model to be processed.

    backend_assignment_fn : Callable[[proto.MIL_pb2.Operation], Optional[List[str]]]
        A function that takes an operation and returns a list of intended backend names.

    Returns
    -------
    MLModel
        A new model instance with updated backend assignments.
    """

    def clone_spec(
        spec: "proto.Model_pb2.Model",
    ) -> "proto.Model_pb2.Model":
        spec_class = spec.__class__
        new_spec = spec_class()
        new_spec.CopyFrom(spec)
        return new_spec

    def set_block_intended_backends(
        block_spec: proto.MIL_pb2.Block,
    ) -> None:
        for operation_spec in block_spec.operations:
            backend_names = backend_assignment_fn(operation_spec)
            set_intended_backends_attr(
                op=operation_spec,
                backend_names=backend_names,
            )
            for operation_block_spec in operation_spec.blocks:
                set_block_intended_backends(block_spec=operation_block_spec)

    model_spec = clone_spec(model.get_spec())

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

    for function_spec in model_spec.mlProgram.functions.values():
        block_spec = function_spec.block_specializations.get(function_spec.opset, None)
        if block_spec is None:
            raise ValueError(
                f"Invalid spec, missing block specialization for opset : {function_spec.opset}"
            )

        set_block_intended_backends(block_spec=block_spec)

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

    Assigns intended backends to operations in the given model.

    This function creates a new MLModel with updated backend assignments for each operation.
    It traverses the entire model structure, applying the provided backend assignment function
    to determine the intended backends for each operation.

    Parameters
    ----------
    model : MLModel
        The input model to be processed.

    backend_assignment_fn : Callable[[proto.MIL_pb2.Operation], Optional[List[str]]]
        A function that takes an operation and returns a list of intended backend names.

    Returns
    -------
    MLModel
        A new model instance with updated backend assignments.
    
coremltools/models/ml_program/experimental/compute_plan_utils.py