def fused_param_init_helper(
    param: torch.Tensor,
    init_fn_: Callable,
    fused_parameters: tuple[int, list[int]],
):
    """Initializes parameters that are fused together.

    Args:
        param (torch.Tensor): Tensor to initialize.
        init_fn_ (Callable): Initialization method.
        fused_parameters (tuple[int, list[int]]): First element of _fused is the dimension
            along which the tensor is fused. Second element is a an iterable of split indices.
    """
    p_ndims = param.ndim
    dim, splits = fused_parameters
    splits = (0, *splits, param.size(dim))  # type: ignore
    for s, e in zip(splits[:-1], splits[1:]):
        # DTensor slicing results in CC and thus produces new Tensor
        # so the update is not inplace, additionally, the init_fn is
        # designed for full tensors, not for a (sharded) DTensor, so
        # we need this context manager to handle the DTensor case
        slice_indices = [slice(None)] * p_ndims  # type: ignore
        slice_indices[dim] = slice(s, e)
        init_fn_(param[slice_indices])  # type: ignore
Initializes parameters that are fused together.

    Args:
        param (torch.Tensor): Tensor to initialize.
        init_fn_ (Callable): Initialization method.
        fused_parameters (tuple[int, list[int]]): First element of _fused is the dimension
            along which the tensor is fused. Second element is a an iterable of split indices.
    
llmfoundry/models/utils/param_init_fns.py