def fused_init_helper_(
    module: nn.Module,
    init_fn_: Callable,
    name_param: str = 'weight',
):
    """Initializes parameters which have been fused for efficiency purposes.

    Parameter initialization is often based on the parameters shape. If a layer is fused,
    initialization should be based on the shapes of the original tensor instead of the
    shape of the fused tensor. Layers which are fused should have the _fused
    attribute. First element of _fused is the dimension along which the tensor is fused.
    Second element is a an iterable of split indices.

    Args:
        module (nn.Module): The module to initialize.
        init_fn_ (Callable): Initialization method.
        name_param (str): Name of parameter to initialize within the module.
    """
    _fused = getattr(module, '_fused', None)
    if _fused is None:
        raise RuntimeError(f'Internal logic error')

    fused_param_init_helper(getattr(module, name_param), init_fn_, _fused)
Initializes parameters which have been fused for efficiency purposes.

    Parameter initialization is often based on the parameters shape. If a layer is fused,
    initialization should be based on the shapes of the original tensor instead of the
    shape of the fused tensor. Layers which are fused should have the _fused
    attribute. First element of _fused is the dimension along which the tensor is fused.
    Second element is a an iterable of split indices.

    Args:
        module (nn.Module): The module to initialize.
        init_fn_ (Callable): Initialization method.
        name_param (str): Name of parameter to initialize within the module.
    
llmfoundry/models/utils/param_init_fns.py