def stacked_init_helper_(
    module: nn.Module,
    init_fn_: Callable,
    name_param: str = 'weight',
):
    """Initializes parameters stacked along a new dimension.

    Parameter initialization is often based on the parameters shape. If a layer is stacked,
    initialization should be based on the shapes of the original tensor instead of the
    shape of the stacked tensor. Layers which are fused should have the _stacked_dim
    attribute defining the new dimension along which they are stacked.

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

    stacked_param_init_helper(getattr(module, name_param), init_fn_, stack_dim)
Initializes parameters stacked along a new dimension.

    Parameter initialization is often based on the parameters shape. If a layer is stacked,
    initialization should be based on the shapes of the original tensor instead of the
    shape of the stacked tensor. Layers which are fused should have the _stacked_dim
    attribute defining the new dimension along which they are stacked.

    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