peft.split_params#

gemma.peft.split_params(

params: dict[str, Any],

) → gemma.peft._tree_utils.SplittedParams[source]

Split a nested tree into 2 trees, one with and without ‘lora’ branches.

Example:

params = {
    'dense': {
        'kernel': w,
        'bias': b,
        'lora': {
            'a': a,
            'b': b,
        },
    },
    'other': other,
}


original, lora = peft.split_params(params)

assert original == {
    'dense': {
        'kernel': w,
        'bias': b,
    },
    'other': other,
}
assert lora == {
    'dense': {
        'lora': {
            'a': a,
            'b': b,
        },
    },
}

Parameters:

params – A nested dictionary representing the input tree containing ‘lora’ branches.

Returns:

(original, lora)

Return type:

A named tuple