def stacked_param_init_helper(
    param: torch.Tensor,
    init_fn_: Callable,
    stack_dim: int,
):
    """Initialize parameters stacked along a new dimension.

    Args:
        param (torch.Tensor): Tensor to initialize.
        init_fn_ (Callable): Initialization method.
        stack_dim (int): Dimension along with parameters are stacked
    """
    p_ndims = param.ndim

    for idx in range(param.size(stack_dim)):
        slice_indices = [slice(None)] * p_ndims  # type: ignore
        slice_indices[stack_dim] = idx  # type: ignore
        init_fn_(param[slice_indices])  # type: ignore
Initialize parameters stacked along a new dimension.

    Args:
        param (torch.Tensor): Tensor to initialize.
        init_fn_ (Callable): Initialization method.
        stack_dim (int): Dimension along with parameters are stacked
    
llmfoundry/models/utils/param_init_fns.py