def get_params(params, from_head_prefix=None):
    """Function extracts relevant parameters into flatten dict from model params,
    appends batch normalization statistics if present"""

    # If Both parameters and batch normalization statistics are present
    if "batch_stats" in params:
        # Extract only parameters for the specified head prefix (if specified) and add batch statistics
        if from_head_prefix is not None:
            extracted_params = flatten_dict(unfreeze(params["params"][from_head_prefix]))
            extracted_params.update(flatten_dict(params["batch_stats"][from_head_prefix]))
        else:
            extracted_params = flatten_dict(unfreeze(params["params"]))
            extracted_params.update(flatten_dict(params["batch_stats"]))

    # Only parameters are present
    else:
        if from_head_prefix is not None:
            extracted_params = flatten_dict(unfreeze(params[from_head_prefix]))
        else:
            extracted_params = flatten_dict(unfreeze(params))

    return extracted_params
Function extracts relevant parameters into flatten dict from model params,
    appends batch normalization statistics if present
dpkd/transformers/tests/test_modeling_flax_common.py