def _flip_fan_mode(init_fn_: Callable):
    """Changes the mode of an init_fn_.

    init_fn_'s "mode" is set to operate on standard torch modules eg torch.nn.Linear.
    If a custom layer transposes its weights before they are allied such that it is
    opposite pytorch's conventions, we must flip the fan mode, from fan_in to fan_out.

    Args:
        init_fn_ (Callable): Initialization method.
    """
    _init_fn_ = deepcopy(init_fn_)
    if 'mode' in _init_fn_.keywords:
        if _init_fn_.keywords['mode'] == 'fan_in':
            _init_fn_.keywords['mode'] = 'fan_out'
        elif _init_fn_.keywords['mode'] == 'fan_out':
            _init_fn_.keywords['mode'] = 'fan_in'
    return _init_fn_
Changes the mode of an init_fn_.

    init_fn_'s "mode" is set to operate on standard torch modules eg torch.nn.Linear.
    If a custom layer transposes its weights before they are allied such that it is
    opposite pytorch's conventions, we must flip the fan mode, from fan_in to fan_out.

    Args:
        init_fn_ (Callable): Initialization method.
    
llmfoundry/models/utils/param_init_fns.py