def freeze_optimizer(
    params: Params,
    optimizer: OptaxOptimizer,
    freeze_fun: Optional[Callable[[Tuple[AnyKey, ...], Array], str]],
) -> OptaxOptimizer:
    """
    Given an optimizer, set its gradient update to zero in correspondence to the frozen parameters specified by the
    freeze function.

    Parameters
    ----------
    params: Params
        Model parameters.
    optimizer: OptaxOptimizer
        An optimizer to freeze.
    freeze_fun: Optional[Callable[[Tuple[AnyKey, ...], Array], str]]
        A freeze function. It takes in input a path and its correspondent array of parameters, and it returns
        "trainable" or "frozen". If "frozen", the gradient of those parameters will be set to zero. If the freeze
        function is None, the optimizer is returned as is.

    Returns
    -------
    OptaxOptimizer
        The same optimizer but frozen where specified by the freeze function.
    """
    if freeze_fun is None:
        return optimizer
    partition_params = path_aware_map(freeze_fun, params)
    all_values_in_labels(tree_leaves(partition_params), ["trainable", "frozen"])
    partition_params = freeze(partition_params)
    partition_optimizers = {"trainable": optimizer, "frozen": set_to_zero()}
    return multi_transform(partition_optimizers, partition_params)

    Given an optimizer, set its gradient update to zero in correspondence to the frozen parameters specified by the
    freeze function.

    Parameters
    ----------
    params: Params
        Model parameters.
    optimizer: OptaxOptimizer
        An optimizer to freeze.
    freeze_fun: Optional[Callable[[Tuple[AnyKey, ...], Array], str]]
        A freeze function. It takes in input a path and its correspondent array of parameters, and it returns
        "trainable" or "frozen". If "frozen", the gradient of those parameters will be set to zero. If the freeze
        function is None, the optimizer is returned as is.

    Returns
    -------
    OptaxOptimizer
        The same optimizer but frozen where specified by the freeze function.
    
fortuna/utils/freeze.py